fix: add retries and handle empty files

Author: PJEstrada

sdk/diffgram/core/diffgram_dataset_iterator.py

@@ -46,8 +46,17 @@ def __validate_file_ids(self):
                 'Some file IDs do not belong to the project. Please provide only files from the same project.')
 
     def get_image_data(self, diffgram_file):
+        MAX_RETRIES = 10
         if hasattr(diffgram_file, 'image'):
-            image = imread(diffgram_file.image.get('url_signed'))
+            for i in range(0, MAX_RETRIES):
+                try:
+                    image = imread(diffgram_file.image.get('url_signed'))
+                    break
+                except Exception as e:
+                    if i < MAX_RETRIES:
+                        continue
+                    else:
+                        raise e
             return image
         else:
             raise Exception('Pytorch datasets only support images. Please provide only file_ids from images')
@@ -70,11 +79,18 @@ def get_file_instances(self, diffgram_file):
             sample['x_max_list'] = x_max_list
             sample['y_min_list'] = y_min_list
             sample['y_max_list'] = y_max_list
+        else:
+            sample['x_min_list'] = []
+            sample['x_max_list'] = []
+            sample['y_min_list'] = []
+            sample['y_max_list'] = []
 
         if 'polygon' in instance_types_in_file:
             has_poly = True
             mask_list = self.extract_masks_from_polygon(instance_list, diffgram_file)
             sample['polygon_mask_list'] = mask_list
+        else:
+            sample['polygon_mask_list'] = []
 
         if len(instance_types_in_file) > 2 and has_boxes and has_boxes:
             raise NotImplementedError(
@@ -83,6 +99,7 @@ def get_file_instances(self, diffgram_file):
 
         label_id_list, label_name_list = self.extract_labels(instance_list)
         sample['label_id_list'] = label_id_list
+        sample['instance_types_in_file'] = instance_types_in_file
         sample['label_name_list'] = label_name_list
 
         return sample

sdk/diffgram/pytorch_diffgram/diffgram_pytorch_dataset.py

@@ -28,7 +28,6 @@ def __get_next_page_of_data(self):
     def __getitem__(self, idx):
         if torch.is_tensor(idx):
             idx = idx.tolist()
-
         diffgram_file = self.project.file.get_by_id(self.diffgram_file_id_list[idx], with_instances = True)
 
         sample = self.get_file_instances(diffgram_file)