Merge pull request #488 from microsoft/PreRelease

Add all other class confidences to the result
microsoft · May 3, 2024 · ef196f5 · ef196f5
2 parents 5f0f09f + 768c68d
commit ef196f5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/PytorchWildlife/models/classification/resnet/amazon.py b/PytorchWildlife/models/classification/resnet/amazon.py
@@ -93,13 +93,16 @@ def results_generation(self, logits, img_ids, id_strip=None):
         probs = torch.softmax(logits, dim=1)
         preds = probs.argmax(dim=1)
         confs = probs.max(dim=1)[0]
+        confidences = probs[0].tolist()
+        result = [[self.CLASS_NAMES[i], confidence] for i, confidence in enumerate(confidences)]
 
         results = []
         for pred, img_id, conf in zip(preds, img_ids, confs):
             r = {"img_id": str(img_id).strip(id_strip)}
             r["prediction"] = self.CLASS_NAMES[pred.item()]
             r["class_id"] = pred.item()
             r["confidence"] = conf.item()
+            r["all_confidences"] = result
             results.append(r)
 
         return results
diff --git a/PytorchWildlife/models/classification/resnet/serengeti.py b/PytorchWildlife/models/classification/resnet/serengeti.py
@@ -67,13 +67,16 @@ def results_generation(self, logits, img_ids, id_strip=None):
         probs = torch.softmax(logits, dim=1)
         preds = probs.argmax(dim=1)
         confs = probs.max(dim=1)[0]
+        confidences = probs[0].tolist()
+        result = [[self.CLASS_NAMES[i], confidence] for i, confidence in enumerate(confidences)]
 
         results = []
         for pred, img_id, conf in zip(preds, img_ids, confs):
             r = {"img_id": str(img_id).strip(id_strip)}
             r["prediction"] = self.CLASS_NAMES[pred.item()]
             r["class_id"] = pred.item()
             r["confidence"] = conf.item()
+            r["all_confidences"] = result
             results.append(r)
 
         return results