Add functionality to get just NMS results (without WBC).

Author: rtgunti

Diff for: predictor.py

@@ -22,6 +22,7 @@
 from multiprocessing import Pool
 import pickle
 import pandas as pd
+from cuda_functions.nms_3D.pth_nms import nms_gpu as nms_3D
 
 
 class Predictor:
@@ -272,7 +273,22 @@ def load_saved_predictions(self, apply_wbc=True):
             pool.close()
             pool.join()
         else:
-            list_of_results_per_patient = list_of_results_per_patient
+            apply_nms = True
+            if apply_nms:
+                self.logger.info('applying NMS to test set predictions with iou = {} and n_ens = {}.'.format(
+                    self.cf.wcs_iou, n_ens))
+                
+                mp_inputs = [[ii[0], ii[1], self.cf.class_dict, self.cf.wcs_iou, n_ens] for ii in list_of_results_per_patient]
+                # pool = Pool(processes=6)
+                # list_of_results_per_patient = pool.map(apply_nms_to_patient, mp_inputs, chunksize=1)
+                # pool.close()
+                # pool.join()
+                list_of_results_per_patient = []
+                for mp_input in mp_inputs:
+                    list_of_results_per_patient.append(apply_nms_to_patient(mp_input))
+                # list_of_results_per_patient = apply_nms_to_patient(mp_inputs[0])
+            else:
+                list_of_results_per_patient = list_of_results_per_patient #returns all the predictions without NMS or WBC
 
         # merge 2D box predictions to 3D cubes (if model predicts 2D but evaluation is run in 3D)
         if self.cf.merge_2D_to_3D_preds:
@@ -501,7 +517,7 @@ def batch_tiling_forward(self, batch):
                 if self.mode == 'val':
                     chunk_dicts += [self.net.train_forward(b, is_validation=True)]
                 else:
-#                     print('check net.test_forward')
+                    # print('check net.test_forward')
                     chunk_dicts += [self.net.test_forward(b, return_masks=self.cf.return_masks_in_test)]
 
 
@@ -520,7 +536,43 @@ def batch_tiling_forward(self, batch):
 
         return results_dict
 
+def apply_nms_to_patient(inputs):
+    '''
+    Apply 3D NMS instead of WBC
+    '''
+    in_patient_results_list, pid, class_dict, wcs_iou, n_ens = inputs
+    out_patient_results_list = [[] for _ in range(len(in_patient_results_list))]
+    for bix, b in enumerate(in_patient_results_list):
+
+        for cl in list(class_dict.keys()):
+
+            boxes = [(ix, box) for ix, box in enumerate(b) if (box['box_type'] == 'det' and box['box_pred_class_id'] == cl)]
+            box_coords = np.array([b[1]['box_coords'] for b in boxes])
+            box_scores = np.array([b[1]['box_score'] for b in boxes])
+            box_center_factor = np.array([b[1]['box_patch_center_factor'] for b in boxes])
+            box_n_overlaps = np.array([b[1]['box_n_overlaps'] for b in boxes])
+            box_patch_id = np.array([b[1]['patch_id'] for b in boxes])
+
+            # print(box_coords.shape, box_scores.shape)
+            # print(box_coords.dtype, box_scores.dtype)
+            box_coords = torch.from_numpy(box_coords).type(torch.float32).cuda()
+            box_scores = torch.from_numpy(box_scores).type(torch.float32).cuda()
+            # print(box_coords.dtype, box_scores.dtype)
+            if 0 not in box_scores.shape:
+                keep_scores, keep_coords = non_max_suppression_3D(
+                    torch.cat((box_coords, box_scores[:, None]), dim=1), box_patch_id, wcs_iou, n_ens)
+                # keep_scores, keep_coords = non_max_suppression_3D(
+                #     np.concatenate((box_coords, box_scores[:, None], box_center_factor[:, None],
+                #                     box_n_overlaps[:, None]), axis=1), box_patch_id, wcs_iou, n_ens)
+
+                for boxix in range(len(keep_scores)):
+                    out_patient_results_list[bix].append({'box_type': 'det', 'box_coords': keep_coords[boxix],
+                                             'box_score': keep_scores[boxix], 'box_pred_class_id': cl})
+
+        # add gt boxes back to new output list.
+        out_patient_results_list[bix].extend([box for box in b if box['box_type'] == 'gt'])
 
+    return [out_patient_results_list, pid]
 
 def apply_wbc_to_patient(inputs):
     """
@@ -603,7 +655,20 @@ def merge_2D_to_3D_preds_per_patient(inputs):
 
     return [out_patient_results_list, pid]
 
+def non_max_suppression_3D(dets, box_patch_id, thresh, n_ens):
+    '''
+    Non max suppression
+    '''
+    # print(dets.shape, box_patch_id.shape)
+    keep = nms_3D(dets, thresh)
+    # print(keep.shape)
+    dets_keep = dets[keep]
+    # print(dets_keep.shape)
+    keep_coords = dets_keep[:,:6].cpu().detach().numpy()
+    keep_scores = dets_keep[:,-1].cpu().detach().numpy()
+    # print(keep_scores.shape, keep_coords.shape)
 
+    return keep_scores, keep_coords
 
 def weighted_box_clustering(dets, box_patch_id, thresh, n_ens):
     """