Update app.py

app.py

@@ -1,58 +1,58 @@
-import cv2
-import gradio as gr
-from huggingface_hub import hf_hub_download
-
-from vision.ssd.mobilenet_v2_ssd_lite import (
-    create_mobilenetv2_ssd_lite,
-    create_mobilenetv2_ssd_lite_predictor,
-)
-
-MODEL_REPO = "fa0311/oita-ken-strawberries-mobilenet"
-MODEL_FILENAME = "20250129_053504/mb2-ssd-lite-Epoch-55-Loss-1.508891262114048.pth"
-LABELS_FILENAME = "20250129_053504/labels.txt"
-
-model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
-label_path = hf_hub_download(repo_id=MODEL_REPO, filename=LABELS_FILENAME)
-
-with open(label_path, "r") as f:
-    class_names = [name.strip() for name in f.readlines()]
-
-net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
-net.load(model_path)
-predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
-
-
-def detect_objects(image, threshold):
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    boxes, labels, probs = predictor.predict(image, 10, threshold)
-
-    for i in range(boxes.size(0)):
-        box = list(map(int, boxes[i, :]))
-        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4)
-        label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
-        cv2.putText(
-            image,
-            label,
-            (box[0] + 10, box[1] + 25),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.8,
-            (255, 0, 255),
-            2,
-        )
-
-    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-
-iface = gr.Interface(
-    fn=detect_objects,
-    inputs=[
-        gr.Image(type="numpy"),
-        gr.Slider(0.1, 1.0, value=0.7, label="Detection Threshold"),
-    ],
-    outputs=gr.Image(type="numpy"),
-    title="SSD Object Detection",
-    description="Upload an image of strawberries to detect objects using MobileNetV2-SSD-Lite.",
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import cv2
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+from vision.ssd.mobilenet_v2_ssd_lite import (
+    create_mobilenetv2_ssd_lite,
+    create_mobilenetv2_ssd_lite_predictor,
+)
+
+MODEL_REPO = "fa0311/oita-ken-strawberries-mobilenet"
+MODEL_FILENAME = "20250129_053504/mb2-ssd-lite-Epoch-55-Loss-1.508891262114048.pth"
+LABELS_FILENAME = "20250129_053504/labels.txt"
+
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
+label_path = hf_hub_download(repo_id=MODEL_REPO, filename=LABELS_FILENAME)
+
+with open(label_path, "r") as f:
+    class_names = [name.strip() for name in f.readlines()]
+
+net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
+net.load(model_path)
+predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
+
+
+def detect_objects(image, threshold):
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    boxes, labels, probs = predictor.predict(image, 10, threshold)
+
+    for i in range(boxes.size(0)):
+        box = list(map(int, boxes[i, :]))
+        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4)
+        label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
+        cv2.putText(
+            image,
+            label,
+            (box[0] + 10, box[1] + 25),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.8,
+            (255, 0, 255),
+            2,
+        )
+
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+
+iface = gr.Interface(
+    fn=detect_objects,
+    inputs=[
+        gr.Image(type="numpy"),
+        gr.Slider(0.1, 1.0, value=0.7, label="Detection Threshold"),
+    ],
+    outputs=gr.Image(type="numpy"),
+    title="SSD Object Detection - Strawberry quality classification",
+    description="Upload an image of strawberries to detect objects using MobileNetV2-SSD-Lite.",
+)
+
+if __name__ == "__main__":
+    iface.launch()