From 6a1ad027549b2e628e548ddc7731c41af8464140 Mon Sep 17 00:00:00 2001 From: LashaO Date: Thu, 13 Jun 2024 18:26:30 +0400 Subject: [PATCH 1/3] Adds multilabel efficientnet --- wbia/algo/detect/efficientnet.py | 80 ++++++++++++++++++++++++++++++-- 1 file changed, 77 insertions(+), 3 deletions(-) diff --git a/wbia/algo/detect/efficientnet.py b/wbia/algo/detect/efficientnet.py index 238599411..57ab5ceb9 100644 --- a/wbia/algo/detect/efficientnet.py +++ b/wbia/algo/detect/efficientnet.py @@ -30,7 +30,8 @@ 'deer_effnet_v0': 'https://wildbookiarepository.azureedge.net/models/labeler_deer_effnet.v0.zip', 'leopard_shark_effnet_v0': 'https://wildbookiarepository.azureedge.net/models/labeler_leopard_shark_effnet.v0.zip', 'trout_effnet_v0': 'https://wildbookiarepository.azureedge.net/models/labeler_trout_effnet.v0.zip', - 'shark_effnet_v0': 'https://wildbookiarepository.azureedge.net/models/labeler_shark_effnet.v0.zip' + 'shark_effnet_v0': 'https://wildbookiarepository.azureedge.net/models/labeler_shark_effnet.v0.zip', + 'msv2_multilabel_effnet_v0': 'https://cthulhu.dyn.wildme.io/public/models/labeler_msv2_multilabel_effnet.v2.zip', } @@ -126,13 +127,49 @@ def _init_transforms(**kwargs): class EfficientnetModel(nn.Module): - def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False): + def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False, multilabel=False): super().__init__() self.model = timm.create_model(model_arch, pretrained=pretrained) + self.multilabel = multilabel + + self.labels = np.array(['back', 'down', 'front', 'left', 'right', 'up']) + self.sort_weights = { + 'up': 1, + 'down': 1, + 'front': 2, + 'back': 2, + 'left': 3, + 'right': 3, + } + + self.reverse_label_map = { + 'up':1, + 'down':2, + 'front':3, + 'back':4, + 'left':5, + 'right':6, + 'upfront':7, + 'upback':8, + 'upleft':9, + 'upright':10, + 'downfront':11, + 'downback':12, + 'downleft':13, + 'downright':14, + 'frontleft':15, + 'frontright':16, + 'backleft':17, + 'backright':18, + } + + if multilabel: + n_class = len(self.sort_weights) if n_class is not None: n_features = self.model.classifier.in_features self.model.classifier = nn.Linear(n_features, n_class) + else: self.model.classifier = nn.Identity(n_features, n_class) ''' @@ -142,8 +179,44 @@ def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False nn.Linear(n_features, n_class, bias=True) ) ''' + + def process_row(self, row_labels, preds, sort_weights, labels): + multi_labels = labels[row_labels == True] + preds = preds[row_labels == True] + # Combine the multi_labels and preds into a list of tuples + label_pred_weight = [(label, pred, sort_weights[label]) for label, pred in zip(multi_labels, preds)] + # Sort by weights first, then by prediction values in descending order + label_pred_weight.sort(key=lambda x: (x[2], -x[1])) + # Create a dictionary to keep the highest value string for each weight + best_labels = {} + for label, pred, weight in label_pred_weight: + if weight not in best_labels or pred > best_labels[weight][1]: + best_labels[weight] = (label, pred) + # Extract the labels in the order of weights - top 2 labels + sorted_labels = [best_labels[weight][0] for weight in sorted(best_labels)[:2]] + return sorted_labels + + def process_multilabel_preds(self, image_preds, sort_weights, labels, reverse_label_map): + multi_label_matrix = (image_preds > 0.5).cpu().numpy() + sorted_labels = [self.process_row(row, preds, sort_weights, labels) for row, preds in zip(multi_label_matrix, image_preds)] + fused_labels = [''.join(x) for x in sorted_labels] + + num_labels = len(reverse_label_map) + one_hot_matrix = np.zeros((len(fused_labels), num_labels)) + + for i, label in enumerate(fused_labels): + if label in reverse_label_map: + one_hot_matrix[i, reverse_label_map[label] - 1] = 1 + + one_hot_tensor = torch.tensor(one_hot_matrix, dtype=torch.float32) + + return one_hot_tensor + def forward(self, x): x = self.model(x) + if self.multilabel: + x = self.process_multilabel_preds(x, self.sort_weights, self.labels, self.reverse_label_map) + return x @@ -553,7 +626,8 @@ def test_single(filepath_list, weights_path, batch_size=1792, multi=PARALLEL, ** num_classes = len(classes) # Initialize the model for this run - model = EfficientnetModel(n_class=num_classes) + multilabel = 'multilabel' in weights_path + model = EfficientnetModel(n_class=num_classes, multilabel=multilabel) # num_ftrs = model.classifier.in_features # model.classifier = nn.Linear(num_ftrs, num_classes) From 5b2322a682b3fb3149a8fc67dd144ec318d7a949 Mon Sep 17 00:00:00 2001 From: LashaO Date: Thu, 13 Jun 2024 19:05:42 +0400 Subject: [PATCH 2/3] fixes linting --- wbia/algo/detect/efficientnet.py | 48 ++++++++++++++++---------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/wbia/algo/detect/efficientnet.py b/wbia/algo/detect/efficientnet.py index 57ab5ceb9..42d2c476d 100644 --- a/wbia/algo/detect/efficientnet.py +++ b/wbia/algo/detect/efficientnet.py @@ -141,26 +141,26 @@ def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False 'left': 3, 'right': 3, } - + self.reverse_label_map = { - 'up':1, - 'down':2, - 'front':3, - 'back':4, - 'left':5, - 'right':6, - 'upfront':7, - 'upback':8, - 'upleft':9, - 'upright':10, - 'downfront':11, - 'downback':12, - 'downleft':13, - 'downright':14, - 'frontleft':15, - 'frontright':16, - 'backleft':17, - 'backright':18, + 'up': 1, + 'down': 2, + 'front': 3, + 'back': 4, + 'left': 5, + 'right': 6, + 'upfront': 7, + 'upback': 8, + 'upleft': 9, + 'upright': 10, + 'downfront': 11, + 'downback': 12, + 'downleft': 13, + 'downright': 14, + 'frontleft': 15, + 'frontright': 16, + 'backleft': 17, + 'backright': 18, } if multilabel: @@ -181,8 +181,8 @@ def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False ''' def process_row(self, row_labels, preds, sort_weights, labels): - multi_labels = labels[row_labels == True] - preds = preds[row_labels == True] + multi_labels = labels[row_labels.astype(bool)] + preds = preds[row_labels.astype(bool)] # Combine the multi_labels and preds into a list of tuples label_pred_weight = [(label, pred, sort_weights[label]) for label, pred in zip(multi_labels, preds)] # Sort by weights first, then by prediction values in descending order @@ -195,7 +195,7 @@ def process_row(self, row_labels, preds, sort_weights, labels): # Extract the labels in the order of weights - top 2 labels sorted_labels = [best_labels[weight][0] for weight in sorted(best_labels)[:2]] return sorted_labels - + def process_multilabel_preds(self, image_preds, sort_weights, labels, reverse_label_map): multi_label_matrix = (image_preds > 0.5).cpu().numpy() sorted_labels = [self.process_row(row, preds, sort_weights, labels) for row, preds in zip(multi_label_matrix, image_preds)] @@ -211,7 +211,7 @@ def process_multilabel_preds(self, image_preds, sort_weights, labels, reverse_la one_hot_tensor = torch.tensor(one_hot_matrix, dtype=torch.float32) return one_hot_tensor - + def forward(self, x): x = self.model(x) if self.multilabel: @@ -626,7 +626,7 @@ def test_single(filepath_list, weights_path, batch_size=1792, multi=PARALLEL, ** num_classes = len(classes) # Initialize the model for this run - multilabel = 'multilabel' in weights_path + multilabel = 'multilabel' in weights_path model = EfficientnetModel(n_class=num_classes, multilabel=multilabel) # num_ftrs = model.classifier.in_features # model.classifier = nn.Linear(num_ftrs, num_classes) From d920cce12c870e91d3ff21ab574f961d8555c283 Mon Sep 17 00:00:00 2001 From: LashaO Date: Thu, 13 Jun 2024 19:06:51 +0400 Subject: [PATCH 3/3] fixes linting --- wbia/algo/detect/efficientnet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wbia/algo/detect/efficientnet.py b/wbia/algo/detect/efficientnet.py index 42d2c476d..19ce53e3f 100644 --- a/wbia/algo/detect/efficientnet.py +++ b/wbia/algo/detect/efficientnet.py @@ -169,7 +169,7 @@ def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False if n_class is not None: n_features = self.model.classifier.in_features self.model.classifier = nn.Linear(n_features, n_class) - + else: self.model.classifier = nn.Identity(n_features, n_class) ''' @@ -179,7 +179,7 @@ def __init__(self, n_class, model_arch='tf_efficientnet_b4_ns', pretrained=False nn.Linear(n_features, n_class, bias=True) ) ''' - + def process_row(self, row_labels, preds, sort_weights, labels): multi_labels = labels[row_labels.astype(bool)] preds = preds[row_labels.astype(bool)]