Implement CSM visualization for 10 classes.

CSM.py

@@ -3,7 +3,7 @@ import torch.nn.functional as F
 
 
 def calc_yolo_csms(imgs_and_preds: torch.Tensor,
-                   sign: bool = True,
+                   sign: bool = False,
                    rescale: bool = True) -> torch.Tensor:
     '''
     computes the cosine similarity map for given input images X
@@ -22,19 +22,22 @@ def calc_yolo_csms(imgs_and_preds: torch.Tensor,
 
     csms = []  # saliency maps w.r.t. all possible output classes
     imgs = []
+    clss = []
 
     for tup in imgs_and_preds:
         img, pred, frame, x1, y1, x2, y2 = tup
         if not img.requires_grad:
             img.requires_grad_()
         logit = pred[5:]
-        imgs.append(img)
+        cls = torch.argmax(pred[5:])
 
         # rescale network output to avoid gradient obfuscation
         if rescale:
             logit = logit / torch.max(torch.abs(logit)) * 10
 
         classes = len(logit)
+        # only first ten classes
+        classes = 10
 
         deltas = []
         for c in range(classes):
@@ -62,7 +65,7 @@ def calc_yolo_csms(imgs_and_preds: torch.Tensor,
             csm = torch.matmul(deltas, deltas.transpose(1, 2))
         except Exception as e:
             print("error")
-            raise e
+            # raise e
 
         #  division by zero can lead to NaNs
         if torch.isnan(csm).any():
@@ -70,8 +73,10 @@ def calc_yolo_csms(imgs_and_preds: torch.Tensor,
             print("NaNs in csm")
         else:
             print(f'{deltas.mean()}')
+            imgs.append(img)
             csms.append(csm)
-    return imgs, csms
+            clss.append(cls)
+    return imgs, csms, clss
 
 
 def calc_csm(model: torch.nn.Module,

main.py

@@ -46,13 +46,19 @@ classes = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus",
 PATH = "saved_patches/realcat.jpg"
 PATH = "saved_patches/fatcat.jpg"
 PATH = "saved_patches/smallcat.jpg"
+PATH = "saved_patches/person.jpg"
 PATCH_SIZE = 300
 
 total_variation = TotalVariation()
 
 
-def show(img):
-    plt.imshow(img.detach().cpu())
+def show(imgs):
+    f, axarr = plt.subplots(2, len(imgs))
+    for i in range(len(imgs)):
+        try:
+            axarr[0, i].imshow(imgs[i].detach().cpu())
+        except:
+            pass
     plt.show()
 
 
@@ -175,7 +181,9 @@ def calculate_csms(frame, predictions):
             # show(frame.flip(2)[int(y1):int(y2), int(x1):int(x2), :] / 255.)
             # print("done")
 
-    imgs, csms = CSM.calc_yolo_csms(imgs_and_preds)
+    imgs, csms, cls = CSM.calc_yolo_csms(imgs_and_preds)
+
+    return imgs, csms, cls
 
 
 if __name__ == "__main__":
@@ -283,13 +291,14 @@ if __name__ == "__main__":
                     pass
 
                 # calculate Cosine Similarity Matrix
-                imgs, csms = calculate_csms(frame, raw_results)
-                for i in range(len(imgs)):
-                    show(imgs[i])
-                    show(csms[i])
-
-                iou, pred = get_best_prediction(bounding_box, raw_results, 15)  # get cat
-                # iou, pred = get_best_prediction(bounding_box, raw_results, 0)  # get personal
+                imgs, csms, clss = calculate_csms(frame, raw_results)
+                for i in range(len(csms)):
+                    # show only person predictions
+                    if clss[i] == 0:
+                        show([imgs[i]/255, csms[i].T])
+
+                # iou, pred = get_best_prediction(bounding_box, raw_results, 15)  # get cat
+                iou, pred = get_best_prediction(bounding_box, raw_results, 0)  # get personal
                 # iou, pred = get_best_prediction(bounding_box, raw_results, 12)  # get parking meter
                 # iou, pred = get_best_prediction(bounding_box, raw_results, 11)  # get stop sign
                 # iou, pred = get_best_prediction(bounding_box, raw_results, 8)  # get boat