initial commit

.gitattributes

@@ -31,3 +31,4 @@ 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
+*.png filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.7
+
+RUN apt-get update \
+    && apt-get install ffmpeg libsm6 libxext6 -y
+
+RUN pip install yolov5 tator gradio
+
+COPY . ./
+
+CMD [ "python", "-u", "./tator_inference.py" ]

app.py

@@ -0,0 +1,38 @@
+import glob
+import gradio as gr
+from inference import *
+from PIL import Image
+
+
+def gradio_app(image_path):
+    """A function that send the file to the inference pipeline, and filters
+    some predictions before outputting to gradio interface."""
+
+    predictions = run_inference(image_path)
+
+    out_img = Image.fromarray(predictions.render()[0])
+
+    return out_img
+
+
+title = "UWROV Deepsea Detector"
+description = "Gradio demo for UWROV Deepsea Detector: Developed by Peyton " \
+              "Lee, Neha Nagvekar, and Cassandra Lam as part of the " \
+              "Underwater Remotely Operated Vehicles Team (UWROV) at the " \
+              "University of Washington. Deepsea Detector is built on " \
+              "MBARI's Monterey Bay Benthic Object Detector, which can also " \
+              "be found in FathomNet's Model Zoo. The model is trained on " \
+              "data from NOAA Ocean Exploration and FathomNet, " \
+              "with assistance from WoRMS for organism classification. All " \
+              "the images and associated annotations we used can be found in " \
+              "our Roboflow project. "
+
+examples = glob.glob("images/*.png")
+
+gr.Interface(gradio_app,
+             inputs=[gr.inputs.Image(type="filepath")],
+             outputs=gr.outputs.Image(type="pil"),
+             enable_queue=True,
+             title=title,
+             description=description,
+             examples=examples).launch()

inference.py

@@ -0,0 +1,69 @@
+import cv2
+import glob
+import numpy as np
+import torch
+import yolov5
+from typing import Dict, Tuple, Union, List, Optional
+
+
+# -----------------------------------------------------------------------------
+# Configs
+# -----------------------------------------------------------------------------
+
+model_path = "models/deepsea-detector.pt"
+
+
+# -----------------------------------------------------------------------------
+# YOLOv5 class
+# -----------------------------------------------------------------------------
+
+class YOLO:
+    """Wrapper class for loading and running YOLO model"""
+
+    def __init__(self, model_path: str, device: Optional[str] = None):
+
+        # load model
+        self.model = yolov5.load(model_path, device=device)
+
+    def __call__(
+            self,
+            img: Union[str, np.ndarray],
+            conf_threshold: float = 0.25,
+            iou_threshold: float = 0.45,
+            image_size: int = 720,
+            classes: Optional[List[int]] = None) -> torch.Tensor:
+        self.model.conf = conf_threshold
+        self.model.iou = iou_threshold
+
+        if classes is not None:
+            self.model.classes = classes
+
+        # pylint: disable=not-callable
+        detections = self.model(img, size=image_size)
+
+        return detections
+
+
+def run_inference(image_path):
+    """Helper function to execute the inference."""
+
+    predictions = model(image_path)
+
+    return predictions
+
+
+# -----------------------------------------------------------------------------
+# Model Creation
+# -----------------------------------------------------------------------------
+model = YOLO(model_path, device='cpu')
+
+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 = run_inference(test_image)
+
+    print("Done.")

models/deepsea-detector.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14c5a32e23e068bec7849726065f02da91f5cb57eedd4d28186403bdfe4f66a8
+size 175330621

requirements.txt

@@ -0,0 +1,3 @@
+yolov5
+tator
+gradio

tator_inference.py

@@ -0,0 +1,100 @@
+import os
+import logging
+from tempfile import TemporaryFile
+
+import cv2
+import numpy as np
+from PIL import Image
+
+import tator
+import inference
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+# Read environment variables that are provided from TATOR
+host = os.getenv('HOST')
+token = os.getenv('TOKEN')
+project_id = int(os.getenv('PROJECT_ID'))
+media_ids = [int(id_) for id_ in os.getenv('MEDIA_IDS').split(',')]
+frames_per_inference = int(os.getenv('FRAMES_PER_INFERENCE', 30))
+
+# Set up the TATOR API.
+api = tator.get_api(host, token)
+
+# Iterate through each video.
+for media_id in media_ids:
+
+    # Download video.
+    media = api.get_media(media_id)
+    logger.info(f"Downloading {media.name}...")
+    out_path = f"/tmp/{media.name}"
+    for progress in tator.util.download_media(api, media, out_path):
+        logger.info(f"Download progress: {progress}%")
+
+    # Do inference on each video.
+    logger.info(f"Doing inference on {media.name}...")
+    localizations = []
+    vid = cv2.VideoCapture(out_path)
+    frame_number = 0
+
+    # Read *every* frame from the video, break when at the end.
+    while True:
+        ret, frame = vid.read()
+        if not ret:
+            break
+
+        # Create a temporary file, access the image data, save data to file.
+        framefile = TemporaryFile(suffix='.jpg')
+        im = Image.fromarray(frame)
+        im.save(framefile)
+
+        # For every N frames, make a prediction; append prediction results
+        # to a list, increase the frame count.
+        if frame_number % frames_per_inference == 0:
+
+            spec = {}
+
+            # Predictions contains all information inside pandas dataframe
+            predictions = inference.run_inference(framefile)
+
+            for i, r in predictions.pandas().xyxy[0].iterrows:
+
+                spec['media_id'] = media_id
+                spec['type'] = None # Unsure, docs not specific
+                spec['frame'] = frame_number
+
+                x, y, x2, y2 = r['xmin'], r['ymin'], r['xmax'], r['ymax']
+                w, h = x2 - x, y2 - y
+
+                spec['x'] = x
+                spec['y'] = y
+                spec['width'] = w
+                spec['height'] = h
+                spec['class_category'] = r['name']
+                spec['confidence'] = r['confidence']
+
+                localizations.append(spec)
+
+        frame_number += 1
+
+    # End interaction with video properly.
+    vid.release()
+
+    logger.info(f"Uploading object detections on {media.name}...")
+
+    # Create the localizations in the video.
+    num_created = 0
+    for response in tator.util.chunked_create(api.create_localization_list,
+                                              project_id,
+                                              localization_spec=localizations):
+        num_created += len(response.id)
+
+    # Output pretty logging information.
+    logger.info(f"Successfully created {num_created} localizations on "
+                f"{media.name}!")
+
+    logger.info("-------------------------------------------------")
+
+logger.info(f"Completed inference on {len(media_ids)} files.")