Initial commit

.gitignore

@@ -0,0 +1,2 @@
+.idea
+runs

README.md

@@ -1,6 +1,6 @@
 ---
 title: MarathonBibDetector
-emoji: 📚
+emoji: 🏃
 colorFrom: blue
 colorTo: red
 sdk: gradio
@@ -8,7 +8,5 @@ sdk_version: 6.3.0
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: 'A detector for the bib wore by a marathon''s participants '
+short_description: 'A detector for the bib wore by a marathon''s participants'
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import gradio
+import ultralytics
+import cv2
+import huggingface_hub
+
+
+project_directory = Path(huggingface_hub.snapshot_download(repo_id="Faraphel/MarathonBibDetector"))
+weights_directory = project_directory / "weights"
+
+
+
+DETECTION_COLOR: tuple[int, int, int] = (64, 255, 64)
+
+MODELS: dict[str, Path] = {
+    "nano": weights_directory / "bib-detector-nano.pt",
+    "medium": weights_directory / "bib-detector-medium.pt",
+}
+LOADED_MODELS: dict[str, ultralytics.YOLO] = {}
+
+
+def get_model(model_type: str = "nano") -> ultralytics.YOLO:
+    """
+    Get a YOLO model instance from the given model type.
+    :param model_type: the model type to load
+    :return: the YOLO model instance
+    """
+
+    # check if the model type is valid
+    if model_type not in MODELS:
+        raise ValueError(f"Invalid model type: {model_type}")
+
+    # check if the model is already loaded
+    if model_type in LOADED_MODELS:
+        return LOADED_MODELS[model_type]
+
+    # otherwise load the model
+    LOADED_MODELS[model_type] = ultralytics.YOLO(MODELS[model_type])
+
+    return LOADED_MODELS[model_type]
+
+
+def infer(image, model_type: str = "nano"):
+    model = get_model(model_type)
+
+    # run the model
+    results = model(image)[0]
+
+    # draw the detected bounding boxes
+    for box in results.boxes:
+        x1, y1, x2, y2 = map(int, box.xyxy[0])
+        conf = float(box.conf[0])
+        cls = int(box.cls[0])
+        label = model.names[cls]
+
+        # draw the bounding box
+        cv2.rectangle(image, (x1, y1), (x2, y2), DETECTION_COLOR, 4)
+        # draw the annotation
+        cv2.putText(image, f"{label} {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, DETECTION_COLOR, 2)
+
+    return image
+
+
+application = gradio.Interface(
+    fn=infer,
+    inputs=[
+        gradio.Image(label="Image", type="numpy"),
+        gradio.Dropdown(label="Model Type", choices=list(MODELS.keys()), value="nano")
+    ],
+    outputs="image",
+    flagging_mode="never"
+)
+application.launch()

environment.yml

@@ -0,0 +1,11 @@
+name: marathon-bib-detector
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  - pip
+  - pip:
+      - gradio>=5.49.1
+      - ultralytics>=8.4.0
+      - opencv-python-headless>=4.12.0
+      - huggingface_hub>=0.34.4

runtime.txt

@@ -0,0 +1 @@
+python-3.12