Updated data to hopefully not hang on dataset preprocessing; adjusted predict to use bbox nms

Author: ericryanchan

Diff for: .gitignore

@@ -1,7 +1,8 @@
 __pycache__/
 .ipynb_checkpoints
 *.pyc
-datasets/
-output/
-input/
-models/
+datasets
+output
+input
+models
+*.jpeg

Diff for: data.py

@@ -13,7 +13,8 @@
 import shutil
 from tqdm import tqdm
 from multiprocessing import Process, Queue
-   
+import queue
+
 import pycocotools
 from detectron2.structures import BoxMode
 
@@ -32,20 +33,20 @@ def compute_bbox(mask):
 def to_coco(dataset_path):
     image_paths = [img_path for img_path in absolute_paths(os.path.join(dataset_path, 'images')) if img_path.endswith('.png')]
     #target_paths = absolute_paths(os.path.join(dataset_path, 'targets'))
-    
+
     dataset_dicts = []
-    
+
     for idx, image_path in enumerate(image_paths):
         target_path = os.path.join(dataset_path, 'targets', os.path.splitext(os.path.basename(image_path))[0] + '.pkl')
         with open(target_path, 'rb') as f:
             target = pickle.load(f)
-        
+
         record = {}
         record['file_name'] = image_path
         record['image_id'] = idx
         record['height'] = target['size'][1]
         record['width'] = target['size'][0]
-        
+
         objs = []
         for m in target['masks']:
             annotation = {'segmentation': pycocotools.mask.encode(np.asarray(m, order="F")),
@@ -68,7 +69,7 @@ def get_image_from_url(url):
         response = requests.get(url, stream=True)
         if response.status_code == 200:
             break
-        time.sleep(0.5)
+        time.sleep(2)
     response.raw.decode_content = True
     img = Image.open(BytesIO(response.content)).convert('RGB')
     return img
@@ -77,7 +78,7 @@ def augment(image, masks, crop_size):
         # Brightness
         brightness_factor = np.random.normal()*0.2 + 1
         image = TF.adjust_brightness(image, brightness_factor)
-        
+
         # Contrast
         contrast_factor = np.random.normal()*0.2 + 1
         image = TF.adjust_contrast(image, contrast_factor)
@@ -89,26 +90,26 @@ def augment(image, masks, crop_size):
         translate = np.random.randint(-30, 30, size=2).tolist()
         image = TF.affine(image, angle, translate, scale, shear, resample=PIL.Image.BILINEAR, fillcolor=None)
         masks = [TF.affine(mask, angle, translate, scale, shear, resample=PIL.Image.BILINEAR, fillcolor=None) for mask in masks]
-        
+
         # Random crop
         i, j, h, w = transforms.RandomCrop.get_params(
             image, output_size=(crop_size, crop_size))
-        
+
         image = TF.crop(image, i, j, h, w)
         masks = [TF.crop(mask, i, j, h, w) for mask in masks]
         # Random horizontal flipping
         if np.random.random() > 0.5:
             image = TF.hflip(image)
             masks = [TF.hflip(mask) for mask in masks]
-    
+
         # Random vertical flipping
         if np.random.random() > 0.5:
             image = TF.vflip(image)
             masks = [TF.vflip(mask) for mask in masks]
-        
+
         # squeeze and binarize
         masks = [(np.array(mask)[:, :, 0] > 0.5).astype(np.uint8) for mask in masks]
-        
+
         # prune masks that have no object or only a sliver of an object
         masks = [mask for mask in masks if mask[10:-10, 10:-10].any()]
         return image, masks
@@ -122,76 +123,90 @@ def __init__(self, task_queue, result_queue, img, masks, out_path, crop_size):
         self.masks = masks
         self.out_path = out_path
         self.crop_size = crop_size
-        
+
     def run(self):
         proc_name = self.name
-        while True:
-            index = self.task_queue.get()
-            if index == -1: break
-            sub_img, sub_masks = augment(self.img, self.masks, self.crop_size)
-            target = {'masks': sub_masks, 'size': sub_img.size}
-            save_mask_target(sub_img, target, f'{index:05d}', dataset_path=self.out_path)            
-            self.result_queue.put(index)
-        return
-    
+        while True:#not stopping.is_set():
+            try:
+                index = self.task_queue.get(True, 1)
+                sub_img, sub_masks = augment(self.img, self.masks, self.crop_size)
+                target = {'masks': sub_masks, 'size': sub_img.size}
+                save_mask_target(sub_img, target, f'{index:05d}', dataset_path=self.out_path)
+                self.result_queue.put(index)
+            except queue.Empty:
+                return
+
 def download_dataset(json_path, out_path, samples_per_img=100, num_threads=16, num_processes=4, selected_ids=None, crop_size=256):
 
     if os.path.exists(out_path):
         shutil.rmtree(out_path)
     os.makedirs(os.path.join(out_path, 'images'))
     os.makedirs(os.path.join(out_path, 'targets'))
-    
-    
+
+
     total_images = 0
     with open(json_path) as f:
         data = json.load(f)
-    
+
     if selected_ids is not None:
         # Filter only selected images
         data = [img_obj for img_obj in data if img_obj['External ID'] in selected_ids]
 
     task_queue = Queue()
     result_queue = Queue()
-                                
-    
+
     with tqdm(total=len(data)*samples_per_img) as pbar:
         for img_obj in data:
+
             img_url = img_obj['Labeled Data']
+            if 'objects' not in img_obj['Label']:
+                continue
+
             mask_urls = [instance['instanceURI'] for instance in img_obj['Label']['objects']]
 
+
             img = get_image_from_url(img_url)
+
             masks = list(ThreadPool(num_threads).imap_unordered(get_image_from_url, mask_urls))
-            
+
+            for _ in range(samples_per_img):
+                task_queue.put(total_images)
+                total_images += 1
+
             workers = []
             for proc_index in range(num_processes):
                 p = Worker(task_queue, result_queue, img, masks, out_path, crop_size)
                 p.daemon = True
                 p.start()
                 workers.append(p)
-            
-            for _ in range(samples_per_img):
-                task_queue.put(total_images)
-                total_images += 1
-                            
-            for index in range(samples_per_img):
-                i = result_queue.get()
-                pbar.update(1)
-                
-            for index in range(num_processes):
-                task_queue.put(-1)
+
+
+
             for worker in workers:
-                worker.join()
+                worker.join(200)
+            for worker in workers:
+                if worker.is_alive():
+                    print("Process timed out")
+
+            pbar.update(samples_per_img)
+            # for index in range(samples_per_img):
+            # while True:
+            #     try:
+            #         i = result_queue.get(True, 10)
+            #         pbar.update(1)
+            #     except queue.Empty:
+            #         break
+
 
-                
 def main():
     ##########################
-    json_path = 'datasets/dataset_export_2020-08-18.json'
-    samples_per_img = 1 #300
+    json_path = 'datasets/export-2020-08-21T20 16 28.026Z.json'
+    samples_per_img = 300
     crop_size = 256
     ##########################
     print('download dataset')
-    
-    
+
+
     train_dataset = [
                      'image_part_001.jpg',
                      'image_part_002.jpg',
@@ -201,34 +216,34 @@ def main():
                      'image_part_006.jpg',
                      'image_part_007.jpg',
                      'image_part_008.jpg',
-                     
-                     #'MC171180.JPG',
-                     #'MC171177.JPG',
-                     #'MC171179.JPG',
-                     #'MC171181.JPG',
-                     #'MC171178.JPG',
-        
+
+                     'MC171180.JPG',
+                     'MC171177.JPG',
+                     'MC171179.JPG',
+                     'MC171181.JPG',
+                     'MC171178.JPG',
+
                      '1758_part_004.jpg',
                      '1758_part_011.jpg',
                      '1758_part_008.jpg',
                      '1755_part_004.jpg',
                      '1730_part_009.jpg',
                      '0467_part_010.jpg',
                      '0467_part_003.jpg',
-                     #'1773_train.JPG',
-                     #'1133_train.JPG',
-                     #'0576_train.JPG',
+                     '1773_train.JPG',
+                     '1133_train.JPG',
+                     '0576_train.JPG',
                     ]
-    
-    
-    
+
+
+
     download_dataset(json_path,
                      'datasets/cells_train_256',
                      samples_per_img=samples_per_img,
                      selected_ids=train_dataset,
                      crop_size=crop_size,
-                     num_processes = 1,
-                     num_threads = 7)
+                     num_processes = 12,
+                     num_threads = 16)
 
 if __name__ == '__main__':
-    main()
+    main()