feat: mvp of the space

.gitattributes

@@ -32,4 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+runs/

.python-version

@@ -0,0 +1 @@
+3.10.12

README.md

@@ -1,13 +1,14 @@
 ---
 title: Salmon Vision
-emoji: 🚀
+emoji: 🐟
 colorFrom: red
-colorTo: red
+colorTo: blue
+python_version: 3.10.12
 sdk: gradio
 sdk_version: 5.5.0
 app_file: app.py
 pinned: false
-short_description: wild salmon migration monitoring
+short_description: Wild salmon migration monitoring
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,169 @@
+"""
+Gradio app to showcase the pyronear model for early forest fire detection.
+"""
+
+from collections import Counter
+from pathlib import Path
+from typing import Any, Tuple
+
+import gradio as gr
+import numpy as np
+from ultralytics import YOLO
+
+
+def bgr_to_rgb(a: np.ndarray) -> np.ndarray:
+    """
+    Turn a BGR numpy array into a RGB numpy array when the array `a` represents
+    an image.
+    """
+    return a[:, :, ::-1]
+
+
+def analyze_predictions(yolo_predictions) -> dict[str, Any]:
+    """
+    Analyze the raw `yolo_predictions` and outputs a dict containg information.
+
+    Args:
+        yolo_predictions: result of calling model.track() on a video
+
+    Returns:
+        counts (int): number of distinct identifiers.
+        ids (set[int]): all the assigned identifiers.
+        detected_species (dict[int, int]): mapping from identifier to instance class
+        names (list[str]): the class names used by the model
+    """
+    if len(yolo_predictions) == 0:
+        return {
+            "counts": 0,
+            "ids": set(),
+            "detected_species": {},
+            "names": None,
+        }
+    else:
+        names = yolo_predictions[0].names
+        ids = set()
+        for prediction in yolo_predictions:
+            if prediction.boxes.id:
+                for id in prediction.boxes.id.numpy().astype("int"):
+                    ids.add(id.item())
+        detected_species = {}
+        for id in ids:
+            counter = Counter()
+            for prediction in yolo_predictions:
+                if prediction.boxes.id:
+                    for idd, klass in zip(
+                        prediction.boxes.id.numpy().astype("int"),
+                        prediction.boxes.cls.numpy().astype("int"),
+                    ):
+                        if idd.item() == id:
+                            counter[klass.item()] += 1
+            selected_class = counter.most_common(1)[0][0]
+            detected_species[id] = selected_class
+        return {
+            "counts": len(ids),
+            "ids": ids,
+            "detected_species": detected_species,
+            "names": names,
+        }
+
+
+def prediction_to_str(yolo_predictions) -> str:
+    """
+    Turn the yolo_predictions into a human friendly string.
+    """
+    if len(yolo_predictions) == 0:
+        return "No prediction"
+    else:
+        result = analyze_predictions(yolo_predictions=yolo_predictions)
+        names = result["names"]
+        detected_species = result["detected_species"]
+        ids = result["ids"]
+        summary_str = "\n".join(
+            [
+                f"- The fish with id {id} is a {names.get(klass, 'Unknown')}"
+                for id, klass in detected_species.items()
+            ]
+        )
+        print(summary_str)
+        return f"Detected {len(ids)} salmons in the video clip with ids {ids}:\n{summary_str}"
+
+
+def predict(model: YOLO, video_filepath: Path) -> Tuple[Path, str]:
+    """
+    Main interface function that runs the model on the provided pil_image and
+    returns the exepected tuple to populate the gradio interface.
+
+    Args:
+        model (YOLO): Loaded ultralytics YOLO model.
+        pil_image (PIL): image to run inference on.
+
+    Returns:
+        pil_image_with_prediction (PIL): image with prediction from the model.
+        raw_prediction_str (str): string representing the raw prediction from the
+        model.
+    """
+    project = "runs/track/"
+    name = video_filepath.stem
+    predictions = model.track(
+        source=video_filepath,
+        save=True,
+        tracker="bytetrack.yaml",
+        exist_ok=True,
+        project=project,
+        name=name,
+    )
+    filepath_video_prediction = Path(f"{project}/{name}/{name}.avi")
+    raw_prediction_str = prediction_to_str(yolo_predictions=predictions)
+    return (filepath_video_prediction, raw_prediction_str)
+
+
+def examples(dir_examples: Path) -> list[Path]:
+    """
+    List the images from the dir_examples directory.
+
+    Returns:
+        filepaths (list[Path]): list of image filepaths.
+    """
+    return list(dir_examples.glob("*.mp4"))
+
+
+def load_model(filepath_weights: Path) -> YOLO:
+    """
+    Load the YOLO model given the filepath_weights.
+    """
+    return YOLO(filepath_weights)
+
+
+# Main Gradio interface
+
+MODEL_FILEPATH_WEIGHTS = Path("data/model/weights.pt")
+DIR_EXAMPLES = Path("data/videos/")
+DEFAULT_IMAGE_INDEX = 0
+
+with gr.Blocks() as demo:
+    model = load_model(MODEL_FILEPATH_WEIGHTS)
+    videos_filepaths = examples(dir_examples=DIR_EXAMPLES)
+    print(f"videos_filepaths: {videos_filepaths}")
+    default_value_input = videos_filepaths[DEFAULT_IMAGE_INDEX]
+    input = gr.Video(
+        value=default_value_input,
+        format="mp4",
+        label="input video",
+        sources=["upload"],
+    )
+    output_video = gr.Video(format="mp4", label="model prediction")
+    output_raw = gr.Text(label="raw prediction")
+
+    fn = lambda video_filepath: predict(
+        model=model, video_filepath=Path(video_filepath)
+    )
+    gr.Interface(
+        title="ML model for wild salmon migration monitoring  🐟",
+        fn=fn,
+        inputs=input,
+        outputs=[output_video, output_raw],
+        examples=videos_filepaths,
+        flagging_mode="never",
+    )
+
+demo.launch()

data/model/weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33d5e39e94c3f94badb476743ec9773f5df09b3f2755379f43b3a594fa755bd2
+size 6239129

data/videos/video1-clip.mp4

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

requirements.txt

@@ -0,0 +1,2 @@
+ultralytics==8.3.*
+gradio==5.4.*