SA/counting_utils.py

@@ -5,16 +5,16 @@ import numpy as np
 import os
 
 def gen_counting_label(labels, channel, tag):
-    b, t, _ = labels.size()
+    b, t = labels.size()
     device = labels.device
     counting_labels = torch.zeros((b, channel))
     if tag:
-        ignore = [0, 1, 2, 108, 109, 110, 111, 112, 113, 114]
+        ignore = [0, 1, 107, 108, 109, 110]
     else:
         ignore = []
     for i in range(b):
         for j in range(t):
-            k = labels[i][j][1]
+            k = labels[i][j]
             if k in ignore:
                 continue
             else: