int32 should be int64 in datasets

halo34 · Mar 8, 2023 · f849203 · f849203
1 parent 5981911
commit f849203
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/dataset/static_dataset.py b/dataset/static_dataset.py
@@ -130,7 +130,7 @@ def __getitem__(self, idx):
         indices = [idx, *np.random.randint(self.__len__(), size=additional_objects)]
 
         merged_images = None
-        merged_masks = np.zeros((self.num_frames, 384, 384), dtype=np.int32)
+        merged_masks = np.zeros((self.num_frames, 384, 384), dtype=np.int64)
 
         for i, list_id in enumerate(indices):
             images, masks = self._get_sample(list_id)
@@ -148,8 +148,8 @@ def __getitem__(self, idx):
         target_objects = labels.tolist()
 
         # Generate one-hot ground-truth
-        cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int32)
-        first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int32)
+        cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int64)
+        first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int64)
         for i, l in enumerate(target_objects):
             this_mask = (masks==l)
             cls_gt[this_mask] = i+1

diff --git a/dataset/vos_dataset.py b/dataset/vos_dataset.py
@@ -190,8 +190,8 @@ def __getitem__(self, idx):
         masks = np.stack(masks, 0)
 
         # Generate one-hot ground-truth
-        cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int32)
-        first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int32)
+        cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int64)
+        first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int64)
         for i, l in enumerate(target_objects):
             this_mask = (masks==l)
             cls_gt[this_mask] = i+1