Yolov5 / MD5

camtrapml/detection/models/base.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+from typing import Union
+
+from camtrapml.image.utils import ImageSource
+
+
+class BaseDetector:
+    def __init__(
+        self,
+        model_path: Union[Path, str, None] = None,
+        model_url: str = "",
+        model_hash: str = "",
+        class_map: Union[dict, None] = None,
+    ) -> None:
+        self.model_path = model_path
+        self.model_url = model_url
+        self.model_hash = model_hash
+        self.class_map = class_map
+
+    def load_model(self,) -> None:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        raise NotImplementedError()
+
+    def detect(self, image_source: ImageSource, min_score: float = 0.1) -> list:
+        raise NotImplementedError()
+
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        raise NotImplementedError()

camtrapml/detection/models/megadetector.py

@@ -21,11 +21,51 @@
 from typing import Tuple, Union
 from json import load
 from .tensorflow import TF1ODAPIFrozenModel
+from .yolov5 import YOLOV5
 from ...download import CACHE_HOME
 
 MODEL_FILES_ROUTE = "https://lilablobssc.blob.core.windows.net/models/camera_traps/megadetector/"
 
 
+class MegaDetectorV5a(YOLOV5):
+    """
+    MegaDetector V5a
+    """
+
+    class_map = {
+        1: "animal",
+        2: "human",
+        3: "vehicle",
+    }
+
+    model_name = "megadetector"
+    model_version = "v5a"
+    model_path = CACHE_HOME / "models/megadetector/v5a/md_v5a.0.0.pt"
+    model_url = "https://github.com/microsoft/CameraTraps/releases/download/v5.0/md_v5a.0.0.pt"
+    model_hash = "ec1d7603ec8cf642d6e0cd008ba2be8c"
+
+class MegaDetectorV5a(YOLOV5):
+    """
+    MegaDetector V5a
+    """
+
+    class_map = {
+        1: "animal",
+        2: "human",
+        3: "vehicle",
+    }
+
+    model_name = "megadetector"
+    model_version = "v5a"
+    model_path = CACHE_HOME / "models/megadetector/v5b/md_v5b.0.0.pt"
+    model_url = "https://github.com/microsoft/CameraTraps/releases/download/v5.0/md_v5b.0.0.pt"
+    model_hash = "76749a62006c66b0e34ac4b846ce6942"
+
+
+
+
+
+
 class MegaDetectorV4_1(TF1ODAPIFrozenModel):  # pylint: disable=C0103
     """
     MegaDetector v4.1

camtrapml/detection/models/yolov5.py

@@ -0,0 +1,147 @@
+from pathlib import Path
+from typing import Union
+
+import yolov5
+
+from camtrapml.download import download, hash_file
+from camtrapml.image.utils import ImageSource, load_image
+
+
+class YOLOV5:
+    """
+    YOLOV5
+    """
+
+    class_map = {}
+
+    model_path = None
+    model_version = None
+    model_url = ""
+
+    def __init__(
+        self,
+        model_path: Union[Path, str, None] = None,
+        model_url: str = "",
+        model_hash: str = "",
+        class_map: Union[dict, None] = None,
+    ) -> None:
+        """
+        Initialises a Tensorflow V1 Object Detection Frozen Graph Model.
+
+        Args:
+
+          model_path: Path to the frozen model .pb file.
+        """
+
+        if model_path:
+            self.model_path = model_path
+
+        if model_url:
+            self.model_url = model_url
+
+        if model_hash:
+            self.model_hash = model_hash
+
+        if class_map:
+            self.class_map = class_map
+
+        self._session = None
+        self._tensors = None
+
+        if Path(self.model_path).exists() is False and self.model_url != "":
+            download(self.model_url, self.model_path, "")
+        elif (
+            Path(self.model_path).exists() is True
+            and model_hash in dir(self)
+            and self.model_hash != ""
+        ):
+            local_hash = hash_file(self.model_path)
+            if local_hash != self.model_hash:
+                raise ValueError(
+                    f"Hash mismatch for model {self.model_path}. Local hash: {local_hash}."
+                )
+
+    def load_model(self,) -> None:
+        """
+        Loads the model from the frozen model .pb file.
+        """
+        self._model = yolov5.load(self.model_path)
+
+        self._model.conf = 0.25  # NMS confidence threshold
+        self._model.iou = 0.45  # NMS IoU threshold
+        self._model.agnostic = False  # NMS class-agnostic
+        self._model.multi_label = False  # NMS multiple labels per box
+        self._model.max_det = 1000  # maximum number of detections per image
+
+
+    def close(self) -> None:
+        """
+        Closes the model.
+        """
+        del self._model
+
+    def detect(self, image_source: ImageSource, min_score: float = 0.1) -> list:
+        """
+        Detects objects in the image.
+
+        Args:
+
+          image_source: PIL Image, Path to image or URL to image.
+          min_score: Minimum score to consider a detection.
+
+        Returns:
+
+          A list of tuples containing the bounding box coordinates and the class
+          of the detected object.
+        """
+
+        # Load the detection graph once.
+        if self._model is None:
+            self.load_model()
+
+        # Load the image.
+        results = self._model(image_source).tolist()[0]
+
+        # print(dir())
+
+        ouputs = []
+
+        for detection in results.xywhn:
+            detection = detection.tolist()[0]
+            print(detection)
+            bbox = detection[:4]
+            conf = detection[4]
+            class_id = int(detection[5])
+
+            if conf < min_score:
+                continue
+
+            ouputs.append({
+                "bbox": bbox,
+                "category": class_id,
+                "conf": conf,
+            })
+
+        return ouputs
+
+
+
+    def __enter__(self):
+        self.load_model()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+# model = YOLOV5(
+#     model_path=Path("/home/bencevans/Projects/yolov5/md_v5a.0.0.torchscript"),
+# )
+
+# with model:
+#     boxes = model.detect(
+#         image_source=Path("/pool0/datasets/ena24/ena24/435.jpg"),
+#         min_score=0.1,
+#     )
+
+#     print(boxes)

pyproject.toml

@@ -42,6 +42,7 @@ font-fredoka-one = "^0.0.4"
 requests = "^2"
 matplotlib = "^3"
 scikit-learn = "1.0.1"
+torch = "^1.12.0"
 
 
 [tool.poetry.dev-dependencies]