From ab442f5e74cb0a99fc02555e97005cdbdc61cabf Mon Sep 17 00:00:00 2001
From: Lachlan Kermode <lachlan_kermode@brown.edu>
Date: Sun, 7 May 2023 17:04:20 -0400
Subject: [PATCH] WIP: yolov5 analyser

---
 examples/yolov5-test.yaml                   | 26 +++++------
 src/lib/analysers/Yolov5/core.py            | 51 +++++++++++++++++++++
 src/lib/analysers/Yolov5/info.yaml          |  6 +++
 src/lib/analysers/Yolov5/partial.Dockerfile |  2 +
 whitelist.txt                               |  3 ++
 5 files changed, 75 insertions(+), 13 deletions(-)
 create mode 100644 src/lib/analysers/Yolov5/core.py
 create mode 100644 src/lib/analysers/Yolov5/info.yaml
 create mode 100644 src/lib/analysers/Yolov5/partial.Dockerfile
 create mode 100644 whitelist.txt

diff --git a/examples/yolov5-test.yaml b/examples/yolov5-test.yaml
index 51f9c7c..838f7e3 100644
--- a/examples/yolov5-test.yaml
+++ b/examples/yolov5-test.yaml
@@ -1,15 +1,15 @@
 folder: media/yolov5
-select:
-  name: Local
-  config:
-    source: data/images
-    aggregate: true
+# select:
+#   name: Youtube
+#   config:
+#     in_parallel: false
+#     search_term: Tear gas
+#     uploaded_before: "2015-10-02T00:00:00Z"
+#     uploaded_after: "2015-10-01T00:00:00Z"
+elements_in:
+  - Youtube
 analyse:
-  name: TorchHub
-  config:
-    dev: true
-    repo: ultralytics/yolov5
-    args:
-      - yolov5s
-    kwargs:
-      pretrained: true
+  - name: Frames
+  - name: Yolov5
+    config:
+      in_parallel: false
diff --git a/src/lib/analysers/Yolov5/core.py b/src/lib/analysers/Yolov5/core.py
new file mode 100644
index 0000000..916f845
--- /dev/null
+++ b/src/lib/analysers/Yolov5/core.py
@@ -0,0 +1,51 @@
+import numpy as np
+from lib.common.exceptions import InvalidAnalyserConfigError
+from lib.common.analyser import Analyser
+from lib.common.etypes import Etype, Union, Array
+from lib.util.cvjson import generate_meta
+from lib.etypes.cvjson import CvJson
+
+import torch
+
+class Yolov5(Analyser):
+    in_etype = Union(Array(Etype.Image), Etype.Json)
+    out_etype = CvJson
+    """ Override to always run serially. Otherwise it hangs, presumably due to
+    the parallelisation that tensorflow does under the hood. """
+
+    @property
+    def in_parallel(self):
+        return False
+
+    def pre_analyse(self, config):
+        # TODO: work out how to load in custom weights, rather than using this default.
+        self.model = torch.hub.load("ultralytics/yolov5", "yolov5n")
+
+        # This function assumes that `element.paths` represents an array of images
+        # to be interpreted. The `get_preds` function operates on a single image,
+        # accepting one argument that is a path to an image. It returns a list of
+        # tuples `('classname', 0.8)`, where `'classname'` is a string
+        # representing the class predicted, and `0.8` is the normalized prediction
+        # probability between 0 and 1. See KerasPretrained/core.py in analysers
+        # for an example. """
+        def get_preds(img_path: Path):
+            results = self.model(img_path)
+            # TODO: work out how to return not the `results` object from yolov5,
+            # but a list of tuples as specified above
+            return []
+
+        self.get_preds = get_preds
+
+
+    def analyse_element(self, element, _):
+        self.logger(f"Running inference on frames in {element.id}...")
+        val = Etype.CvJson.from_preds(element, self.get_preds)
+        self.logger("Tried to show results.")
+        self.disk.delete_local_on_write = True
+        return val
+
+    def post_analyse(self, elements) -> Etype.Json.as_array():
+        return generate_meta(elements, logger=self.logger)
+
+
+module = Yolov5
diff --git a/src/lib/analysers/Yolov5/info.yaml b/src/lib/analysers/Yolov5/info.yaml
new file mode 100644
index 0000000..509105b
--- /dev/null
+++ b/src/lib/analysers/Yolov5/info.yaml
@@ -0,0 +1,6 @@
+desc: Classify objects in images using YoloV5 - https://github.com/ultralytics/yolov5. This analyser allows you to run inference using one of the pretrained checkpoints, or to provide custom weights.
+args: []
+  # - name: checkpoint
+  #   desc: The checkpoint you want to use for weights. This can be a local path to custom weights, or one of the available pretrained checkpoints (see https://github.com/ultralytics/yolov5#pretrained-checkpoints).
+  #   required: true
+  #   input: string
diff --git a/src/lib/analysers/Yolov5/partial.Dockerfile b/src/lib/analysers/Yolov5/partial.Dockerfile
new file mode 100644
index 0000000..39c7611
--- /dev/null
+++ b/src/lib/analysers/Yolov5/partial.Dockerfile
@@ -0,0 +1,2 @@
+RUN git clone https://github.com/ultralytics/yolov5
+RUN cd yolov5 && pip install -r requirements.txt
diff --git a/whitelist.txt b/whitelist.txt
new file mode 100644
index 0000000..8be3cfd
--- /dev/null
+++ b/whitelist.txt
@@ -0,0 +1,3 @@
+Youtube
+Frames
+Yolov5