use seavision

.gitignore

@@ -2,6 +2,7 @@
 flagged*/
 gradio_cached_examples/
 *ipynb
+.gradio/
 
 # macOS
 .DS_Store

app.py

@@ -6,6 +6,7 @@ Any new model should implement the following functions:
 """
 import os
 import glob
+import spaces
 import gradio as gr
 from huggingface_hub import get_token
 from utils import (
@@ -15,7 +16,8 @@ from utils import (
     FlaggedCounter,
 )
 from flagging import HuggingFaceDatasetSaver
-from model_yolov5 import load_model, inference
+import install_private_repos
+from seavision import load_model, AHOY
 
 
 TITLE = """
@@ -42,11 +44,13 @@ h1 {
 }
 """
 
-model = load_model("SEA-AI/yolov5n6-RGB", img_size=1280)
-model.conf = 0.10
-model.iou = 0.4
-model.max_det = 100
-model.agnostic = True  # NMS class-agnostic
+model = load_model("ahoy-RGB-b2")
+
+@spaces.GPU
+def inference(model: AHOY, image):
+    """Run inference on image and return annotated image."""
+    results = model(image)
+    return results.draw(image, diameter=4)
 
 # Flagging
 dataset_name = "SEA-AI/crowdsourced-sea-images"

install_private_repos.py

@@ -0,0 +1,10 @@
+import os
+
+GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")
+if GITHUB_TOKEN is not None:
+    try:
+        cmd = f"pip3 install git+https://{GITHUB_TOKEN}@github.com/SEA-AI/seavision.git"
+        os.system(cmd)
+    except Exception as e:
+        print(f"Error: {e}")
+        raise e

model_yolov5.py

@@ -1,25 +0,0 @@
-import spaces
-import numpy as np
-import yolov5
-from yolov5.utils.plots import Annotator, colors
-from huggingface_hub import get_token
-
-
-def load_model(model_path, img_size=640):
-    """Load model from HuggingFace Hub."""
-    model = yolov5.load(model_path, hf_token=get_token())
-    model.img_size = img_size  # add img_size attribute
-    return model
-
-
-@spaces.GPU
-def inference(model, image):
-    """Run inference on image and return annotated image."""
-    results = model(image, size=model.img_size)
-    print(f"{results=}")
-    annotator = Annotator(np.asarray(image))
-    for *box, _, cls in reversed(results.pred[0]):
-        # label = f'{model.names[int(cls)]} {conf:.2f}'
-        # print(f'{cls} {conf:.2f} {box}')
-        annotator.box_label(box, "", color=colors(cls, True))
-    return annotator.im