Fixes this issue: RobotLocomotion/pytorch-dense-correspondence#204

Author: gasiortomasz

dense_correspondence/dataset/spartan_dataset_masked.py

@@ -269,7 +269,7 @@ def get_scene_list(self, mode=None):
             scene_list.append(scene_name)
 
         return scene_list
-    
+
     def get_list_of_objects(self):
         """
         Returns a list of object ids
@@ -643,6 +643,16 @@ def get_within_scene_data(self, scene_name, metadata, for_synthetic_multi_object
         image_a_depth_numpy = np.asarray(image_a_depth)
         image_b_depth_numpy = np.asarray(image_b_depth)
 
+        image_a_mask_numpy = np.asarray(image_a_mask)
+        image_b_mask_numpy = np.asarray(image_b_mask)
+        img_size = np.size(image_a_mask_numpy)
+        min_mask_size = 0.01*img_size
+
+        if (np.sum(image_a_mask_numpy) < min_mask_size) or (np.sum(image_b_mask_numpy) < min_mask_size):
+            logging.info("not enough pixels in mask, skipping")
+            image_a_rgb_tensor = self.rgb_image_to_tensor(image_a_rgb)
+            return self.return_empty_data(image_a_rgb_tensor, image_a_rgb_tensor)
+
         if self.sample_matches_only_off_mask:
             correspondence_mask = np.asarray(image_a_mask)
         else:
@@ -1212,7 +1222,7 @@ def get_first_image_index(self, scene_name):
     @property
     def config(self):
         return self._config
-    
+
     @staticmethod
     def merge_single_object_configs(config_list):
         """