FIX: video

detection.py

@@ -70,10 +70,17 @@ def detect_video(frames, processor, clip_model, detection_model):
     
     pred_score = float(detection_model(last_hidden_states)[0][0].cpu().detach().numpy())
     assert 0 <= pred_score <= 1
-
-    attention_maps = outputs.attentions[-1].cpu().detach().numpy()
-    cls_attention_map = attention_maps[:, :, 0, :]
-    cls_attention_map = cls_attention_map.mean(axis=0)
+    
+    for layer_idx in range(len(outputs.attentions)):
+        attn_map = outputs.attentions[layer_idx]
+        if layer_idx == 0:
+            last_layer_attn = attn_map
+        else:
+            if layer_idx < 6:
+                last_layer_attn += attn_map
+    
+    head_mean_attn = last_layer_attn.mean(dim=1)[0]
+    cls_attention_map = head_mean_attn[0, 1:]
 
     blended_image = vis_attn(image, cls_attention_map)