def test_keep_largest_object(nii_seg): # Set a voxel to 1 at the corner to make sure it is set to 0 by the function coord = (1, 1, 1) nii_seg.dataobj[coord] = 1 # Test function with array input arr_seg_proc = imed_postpro.keep_largest_object(np.copy(np.asanyarray(nii_seg.dataobj))) assert isinstance(arr_seg_proc, np.ndarray) assert check_bin_vs_soft(nii_seg.dataobj, arr_seg_proc) assert arr_seg_proc[coord] == 0 # Make sure it works with nibabel input nii_seg_proc = imed_postpro.keep_largest_object(nii_seg) assert isinstance(nii_seg_proc, nib.nifti1.Nifti1Image) assert check_bin_vs_soft(nii_seg.dataobj, nii_seg_proc.dataobj) assert nii_seg_proc.dataobj[coord] == 0
def generate_bounding_box_file(subject_list, model_path, path_output, gpu_id=0, slice_axis=0, contrast_lst=None, keep_largest_only=True, safety_factor=None): """Creates json file containing the bounding box dimension for each images. The file has the following format: {"path/to/img.nii.gz": [[x1_min, x1_max, y1_min, y1_max, z1_min, z1_max], [x2_min, x2_max, y2_min, y2_max, z2_min, z2_max]]} where each list represents the coordinates of an object on the image (2 instance of a given object in this example). Args: subject_list (list): List of all subjects in the BIDS directory. model_path (string): Path to object detection model. path_output (string): Output path. gpu_id (int): If available, GPU number. slice_axis (int): Slice axis (0: sagittal, 1: coronal, 2: axial). contrast_lst (list): Contrasts. keep_largest_only (bool): Boolean representing if only the largest object of the prediction is kept. safety_factor (list or tuple): Factors to multiply each dimension of the bounding box. Returns: dict: Dictionary containing bounding boxes related to their image. """ bounding_box_dict = {} if safety_factor is None: safety_factor = [1.0, 1.0, 1.0] for subject in subject_list: if subject.record["modality"] in contrast_lst: subject_path = str(subject.record["absolute_path"]) object_mask, _ = imed_inference.segment_volume(model_path, [subject_path], gpu_id=gpu_id) object_mask = object_mask[0] if keep_largest_only: object_mask = imed_postpro.keep_largest_object(object_mask) mask_path = os.path.join(path_output, "detection_mask") if not os.path.exists(mask_path): os.mkdir(mask_path) nib.save(object_mask, os.path.join(mask_path, subject_path.split("/")[-1])) ras_orientation = nib.as_closest_canonical(object_mask) hwd_orientation = imed_loader_utils.orient_img_hwd( ras_orientation.get_fdata(), slice_axis) bounding_boxes = get_bounding_boxes(hwd_orientation) bounding_box_dict[subject_path] = [ adjust_bb_size(bb, safety_factor) for bb in bounding_boxes ] file_path = os.path.join(path_output, 'bounding_boxes.json') with open(file_path, 'w') as fp: json.dump(bounding_box_dict, fp, indent=4) return bounding_box_dict