import os import glob import numpy as np import detectron2 import torchvision import cv2 import torch from detectron2 import model_zoo from detectron2.data import Metadata from detectron2.structures import BoxMode from detectron2.utils.visualizer import Visualizer from detectron2.config import get_cfg from detectron2.utils.visualizer import ColorMode from detectron2.modeling import build_model import detectron2.data.transforms as T from detectron2.checkpoint import DetectionCheckpointer import matplotlib.pyplot as plt from PIL import Image # ----------------------------------------------------------------------------- # CONFIGS - loaded just the one time when script is first ran to save time. # # This is where you will set all the relevant config file and weight file # variables: # CONFIG_FILE - Training specific config file for fathomnet # WEIGHTS_FILE - Path to the model with fathomnet weights # NMS_THRESH - Set a nms threshold for the all boxes results # SCORE_THRESH - This is where you can set the model score threshold CONFIG_FILE = "fathomnet_config_v2_1280.yaml" WEIGHTS_FILE = "model_final.pth" NMS_THRESH = 0.45 # SCORE_THRESH = 0.3 # # A metadata object that contains metadata on each class category; used with # Detectron for linking predictions to names and for visualizations. fathomnet_metadata = Metadata( name='fathomnet_val', thing_classes=[ 'Anemone', 'Fish', 'Eel', 'Gastropod', 'Sea star', 'Feather star', 'Sea cucumber', 'Urchin', 'Glass sponge', 'Sea fan', 'Soft coral', 'Sea pen', 'Stony coral', 'Ray', 'Crab', 'Shrimp', 'Squat lobster', 'Flatfish', 'Sea spider', 'Worm'] ) # This is where the model parameters are instantiated. There is a LOT of # nested arguments in these yaml files, and the merging of baseline defaults # plus dataset specific parameters. base_model_path = "COCO-Detection/retinanet_R_50_FPN_3x.yaml" cfg = get_cfg() cfg.MODEL.DEVICE = 'cpu' cfg.merge_from_file(model_zoo.get_config_file(base_model_path)) cfg.merge_from_file(CONFIG_FILE) cfg.MODEL.RETINANET.SCORE_THRESH_TEST = SCORE_THRESH cfg.MODEL.WEIGHTS = WEIGHTS_FILE # Loading of the model weights, but more importantly this is where the model # is actually instantiated as something that can take inputs and provide # outputs. There is a lot of documentation about this, but not much in the # way of straightforward tutorials. model = build_model(cfg) checkpointer = DetectionCheckpointer(model) checkpointer.load(cfg.MODEL.WEIGHTS) model.eval() # Create two augmentations and make a list to iterate over aug1 = T.ResizeShortestEdge(short_edge_length=[cfg.INPUT.MIN_SIZE_TEST], max_size=cfg.INPUT.MAX_SIZE_TEST, sample_style="choice") aug2 = T.ResizeShortestEdge(short_edge_length=[1080], max_size=1980, sample_style="choice") augmentations = [aug1, aug2] # We use a separate NMS layer because initially detectron only does nms intra # class, so we want to do nms on all boxes. post_process_nms = torchvision.ops.nms # ----------------------------------------------------------------------------- def run_inference(test_image): """This function runs through inference pipeline, taking in a single image as input. The image will be opened, augmented, ran through the model, which will output bounding boxes and class categories for each object detected. These are then passed back to the calling function.""" # Load the image, get the height and width. Iterate over each # augmentation: do the augmentation, run the model, perform nms # thresholding, instantiate a useful object for visualizing the outputs. # Saves a list of outputs objects img = cv2.imread(test_image) im_height, im_width, _ = img.shape v_inf = Visualizer(img[:, :, ::-1], metadata=fathomnet_metadata, scale=1.0, instance_mode=ColorMode.SEGMENTATION) insts = [] # iterate over input augmentations (apply resizing) for augmentation in augmentations: im = augmentation.get_transform(img).apply_image(img) # pre-process image by reshaping and converting to tensor # pass to model, which outputs a dict containing info on all detections with torch.no_grad(): im = torch.as_tensor(im.astype("float32").transpose(2, 0, 1)) model_outputs = model([{"image": im, "height": im_height, "width": im_width}])[0] # populate list with all outputs for _ in range(len(model_outputs['instances'])): insts.append(model_outputs['instances'][_]) # Concatenate the model outputs and run NMS thresholding on all output; # instantiate a dummy Instance object to concatenate the instances model_inst = detectron2.structures.instances.Instances([im_height, im_width]) xx = model_inst.cat(insts)[ post_process_nms(model_inst.cat(insts).pred_boxes.tensor, model_inst.cat(insts).scores, NMS_THRESH).to("cpu").tolist()] print(test_image + ' - Number of predictions:', len(xx)) out_inf_raw = v_inf.draw_instance_predictions(xx.to("cpu")) out_pil = Image.fromarray(out_inf_raw.get_image()).convert('RGB') # Converting the predictions as output by Detectron2, to a TATOR format. predictions = convert_predictions(xx, v_inf.metadata.thing_classes) return predictions, out_pil def convert_predictions(xx, thing_classes): """Helper funtion to post-process the predictions made by Detectron2 codebase to work with TATOR input requirements.""" predictions = [] for _ in range(len(xx)): # Obtain the first prediction, instance instance = xx.__getitem__(_) # Map the coordinates to the variables x, y, x2, y2 = map(float, instance.pred_boxes.tensor[0]) w, h = x2 - x, y2 - y # Use class list to get the common name (string); get confidence score. class_category = thing_classes[int(instance.pred_classes[0])] confidence_score = float(instance.scores[0]) # Create a spec dict for TATOR prediction = {'x': x, 'y': y, 'width': w, 'height': h, 'class_category': class_category, 'confidence': confidence_score} predictions.append(prediction) return predictions if __name__ == "__main__": # For demo purposes: run through a couple of test # images and then output the predictions in a folder. test_images = glob.glob("images/*.png") for test_image in test_images: predictions, out_img = run_inference(test_image) print("Done.")