inference.py

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.")