update base_tracker

tracker/base_tracker.py

@@ -9,9 +9,16 @@ import torch
 import yaml
 from model.network import XMem
 from inference.inference_core import InferenceCore
+from inference.data.mask_mapper import MaskMapper
+
 # for data transormation
 from torchvision import transforms
 from dataset.range_transform import im_normalization
+import torch.nn.functional as F
+
+import sys
+sys.path.insert(0, sys.path[0]+"/../")
+from tools.painter import mask_painter
 
 
 class BaseTracker:
@@ -32,6 +39,14 @@ class BaseTracker:
             transforms.ToTensor(),
             im_normalization,
         ])
+        self.device = device
+
+    def resize_mask(self, mask):
+        # mask transform is applied AFTER mapper, so we need to post-process it in eval.py
+        h, w = mask.shape[-2:]
+        min_hw = min(h, w)
+        return F.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)), 
+                    mode='nearest')
 
     def track(self, frames, first_frame_annotation):
         """
@@ -40,14 +55,36 @@ class BaseTracker:
         first_frame_annotation: numpy array: H, W
 
         Output:
-        masks: numpy arrays: T, H, W
+        masks: numpy arrays: H, W
         """
+        shape = np.array(frames).shape[1:3] # H, W
+        frame_list = [self.im_transform(frame).to(self.device) for frame in frames]
+        frame_tensors = torch.stack(frame_list, dim=0)
+        
         # data transformation
-        for frame in frames:
-            frame = self.im_transform(frame)
+        mapper = MaskMapper()
 
-        # tracking
+        vid_length = len(frame_tensors)
 
+        for ti, frame_tensor in enumerate(frame_tensors):
+            if ti == 0:
+                mask, labels = mapper.convert_mask(first_frame_annotation)
+                mask = torch.Tensor(mask).to(self.device)
+                self.tracker.set_all_labels(list(mapper.remappings.values()))
+            else:
+                mask = None
+                labels = None
+            
+            # track one frame
+            prob = self.tracker.step(frame_tensor, mask, labels, end=(ti==vid_length-1))
+            
+            out_mask = torch.argmax(prob, dim=0)
+            out_mask = (out_mask.detach().cpu().numpy()).astype(np.uint8)
+
+            painted_image = mask_painter(frames[ti], out_mask)
+            # save
+            painted_image = Image.fromarray(painted_image)
+            painted_image.save(f'/ssd1/gaomingqi/results/TrackA/{ti}.png')
 
 
 if __name__ == '__main__':
@@ -56,20 +93,17 @@ if __name__ == '__main__':
     video_path_list.sort()
     # first frame
     first_frame_path = '/ssd1/gaomingqi/datasets/davis/Annotations/480p/dance-twirl/00000.png'
-
     # load frames
-<<<<<<< HEAD
-=======
-    frames = ["test_confict"]
->>>>>>> a5606340a199569856ffa1585aeeff5a40cc34ba
+    frames = []
     for video_path in video_path_list:
         frames.append(np.array(Image.open(video_path).convert('RGB')))
     frames = np.stack(frames, 0)    # N, H, W, C
-
     # load first frame annotation
     first_frame_annotation = np.array(Image.open(first_frame_path).convert('P'))    # H, W, C
 
+    # ----------------------------------------------------------
     # initalise tracker
+    # ----------------------------------------------------------
     device = 'cuda:0'
     XMEM_checkpoint = '/ssd1/gaomingqi/checkpoints/XMem-s012.pth'
     tracker = BaseTracker('cuda:0', XMEM_checkpoint)