updated app

.gitattributes

@@ -31,3 +31,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/scene_5.png filter=lfs diff=lfs merge=lfs -text
+images/fish_2.png filter=lfs diff=lfs merge=lfs -text
+images/jelly_2.png filter=lfs diff=lfs merge=lfs -text
+images/scene_2.png filter=lfs diff=lfs merge=lfs -text

inference.py

@@ -0,0 +1,199 @@
+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
+
+# -----------------------------------------------------------------------------
+# 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.IMAGE_BW)
+
+    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.")
+