class DetectronBladder: def __init__(self): self.image_utils = ImageUtils() self.mask_utils = MaskUtils() self.file_utils = FileUtils() self.coco_utils = CocoUtils() self.cfg = None self.predictor = None self.bladder_metadata = None self.configure() def configure(self): self.cfg = get_cfg() self.cfg.OUTPUT_DIR = "./Output/Bladder" self.cfg.merge_from_file( model_zoo.get_config_file( "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, "model_final.pth") self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 self.predictor = DefaultPredictor(self.cfg) self.bladder_metadata = MetadataCatalog.get("bladder_test") def train(self): os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) trainer.train() def register_dataset(self): for d in ["train", "test"]: register_coco_instances(f"bladder_{d}", {}, f"Data/Vezica/{d}.json", f"Data/Vezica/{d}") def process_result(self, image_path, outputs, output): try: instances = outputs["instances"].to("cpu") if instances.has("pred_boxes") and len(instances.pred_boxes) > 0: if output != "": im2 = Image.open(image_path) img = self.image_utils.get_masked_image( im2, instances.pred_masks[0]) img.save(output) return True, instances.pred_masks[:1], "" else: return False, [], "" except Exception as e: return False, [], str(e) def test_detect(self, images_folder, output_folder, test_files, width, height): ious = [] dscs = [] for (image_path, json_path) in test_files: output_path = output_folder + image_path is_bladder, pred_masks, err = self.detect_bladder( images_folder + image_path, output_path) if json_path is not None: masks = self.coco_utils.json_annotation_to_masks( images_folder + json_path) else: masks = None output_mask = self.mask_utils.combine_masks( pred_masks, width, height) self.mask_utils.draw_mask(output_mask, output_path[:-4] + "-m.png") expected_mask = self.mask_utils.combine_masks(masks, width, height) self.mask_utils.draw_mask(expected_mask, output_path[:-4] + "-m2.png") intersection, union, sum = self.mask_utils.get_intersection_union_sum( expected_mask, output_mask) if union == 0: if intersection == 0: iou = 1 else: iou = 0 else: iou = intersection / union if iou <= 0.1: print(image_path) if sum == 0: dsc = 1 else: dsc = (2 * intersection) / sum ious.append(iou) dscs.append(dsc) return ious, dscs def get_IOU_DSC(self, images_folder, output_folder): test_files = self.file_utils.get_test_files(images_folder) ious, dscs = self.test_detect(images_folder, output_folder, test_files, 512, 512) return sum(ious) / len(ious), sum(dscs) / len(dscs) def detect_bladder(self, image_path, output_path): im = cv2.imread(image_path) outputs = self.predictor(im) return self.process_result(image_path, outputs, output_path)
class DetectronBladderTumor: def __init__(self): self.image_utils = ImageUtils() self.mask_utils = MaskUtils() self.file_utils = FileUtils() self.coco_utils = CocoUtils() self.cfg = None self.predictor = None self.tumor_metadata = None self.configure() def configure(self): self.cfg = get_cfg() self.cfg.OUTPUT_DIR = "./Output/Tumor" self.cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")) self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 self.cfg.MODEL.WEIGHTS = os.path.join(self.cfg.OUTPUT_DIR, "model_final.pth") self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 self.predictor = DefaultPredictor(self.cfg) self.tumor_metadata = MetadataCatalog.get("tumor_test") def train(self): os.makedirs(self.cfg.OUTPUT_DIR, exist_ok=True) trainer = DefaultTrainer(self.cfg) trainer.resume_or_load(resume=False) trainer.train() def save_resulted_image(self, im, instances, output_path): v = Visualizer(im[:, :, ::-1], metadata=self.tumor_metadata, scale=0.8, instance_mode=ColorMode.IMAGE_BW ) v = v.draw_instance_predictions(instances) plt.figure(figsize=(14, 10)) img = v.get_image()[:, :, ::-1] plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) if output_path != "": resized = cv2.resize(img, (512, 512), interpolation=cv2.INTER_AREA) cv2.imwrite(output_path, resized) plt.close() def process_result(self, outputs): instances = outputs["instances"].to("cpu") # self.save_resulted_image(im, instances, output_path) if instances.has("pred_masks") and len(instances.pred_masks) > 0: return instances.pred_masks, "" else: return [], "" def detect_tumor(self, image_path): im = cv2.imread(image_path) outputs = self.predictor(im) return self.process_result(outputs) def test_detect(self, images_folder, output_folder, test_files, width, height): ious = [] dscs = [] for (image_path, json_path) in test_files: output_path = output_folder + image_path pred_masks, err = self.detect_tumor(images_folder + image_path) if json_path is not None: masks = self.coco_utils.json_annotation_to_masks(images_folder + json_path) else: masks = None output_mask = self.mask_utils.combine_masks(pred_masks, width, height) self.mask_utils.draw_mask(output_mask, output_path[:-4] + "-m.png") expected_mask = self.mask_utils.combine_masks(masks, width, height) self.mask_utils.draw_mask(expected_mask, output_path[:-4] + "-m2.png") intersection, union, sum = self.mask_utils.get_intersection_union_sum(expected_mask, output_mask) if union == 0: if intersection == 0: iou = 1 else: iou = 0 else: iou = intersection/union if iou <= 0.1: print(image_path) if sum == 0: dsc = 1 else: dsc = (2 * intersection) / sum ious.append(iou) dscs.append(dsc) return ious, dscs def get_IOU_DSC(self, images_folder, output_folder): test_files = self.file_utils.get_test_files(images_folder) ious, dscs = self.test_detect(images_folder, output_folder, test_files, 512, 512) return sum(ious) / len(ious), sum(dscs) / len(dscs)