obj detect only

app.py

@@ -1,259 +1,165 @@
-#Main Gradio app ith image classification and object detection tabs
-import gradio as gr
-import torch
-import torch.nn.functional as F
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-from PIL import Image
+# app.py — Object Detection only (multi-image YOLO, up to 10)
 import os
 import csv
 import tempfile
 from pathlib import Path
-from ultralytics import YOLO
-# ultralytics YOLO import (for object detection)
+from typing import List, Tuple
+
+import gradio as gr
+from PIL import Image
+
+# Try import ultralytics (ensure it's in requirements.txt)
 try:
     from ultralytics import YOLO
 except Exception:
     YOLO = None
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-MODEL_ID = "dshi01/convnext-tiny-224-7clss"
-
-print(f"Loading model from: {MODEL_ID}")
-processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
-model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
-model.eval()
-
-# (Optional) use model's own labels if present
-ID2LABEL = [
-    model.config.id2label.get(str(i), model.config.id2label.get(i, f"Label_{i}"))
-    for i in range(model.config.num_labels)
-]
-def classify_image(image):
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image).convert("RGB")
-
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = F.softmax(logits, dim=1)[0].tolist()
-
-    return {ID2LABEL[i]: float(p) for i, p in enumerate(probs)}
-
-# ---------- NEW: batch classify up to 10 images ----------
 MAX_BATCH = 10
 
-def classify_images_batch(files):
-    """
-    files: list of gradio UploadedFile (paths) or None
-    Returns:
-      - gallery: list of (image, caption)
-      - table: list of rows for Dataframe
-    """
-    if not files:
-        return [], [], None
-
-    # Keep at most 10
-    files = files[:MAX_BATCH]
-
-    # Load as PIL
-    pil_images, names = [], []
-    for f in files:
-        path = getattr(f, "name", None) or getattr(f, "path", None) or f
-        try:
-            img = Image.open(path).convert("RGB")
-            pil_images.append(img)
-            names.append(os.path.basename(path))
-        except Exception:
-            # Skip unreadable file
-            continue
-
-    if not pil_images:
-        return [], [], None
-
-    # Batch preprocess + forward
-    inputs = processor(images=pil_images, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = F.softmax(logits, dim=1)
-
-    # Build outputs
-    gallery = []
-    table_rows = []  # [filename, top1_label, top1_conf, top3_labels, top3_confs]
-
-    for idx, (img, fname) in enumerate(zip(pil_images, names)):
-        p = probs[idx].tolist()
-        top_idxs = sorted(range(len(p)), key=lambda i: p[i], reverse=True)[:3]
-        top1 = top_idxs[0]
-        caption = f"{ID2LABEL[top1]} ({p[top1]:.2%})"
+# Option A: local file baked into Space (easiest if allowed)
+YOLO_WEIGHTS = os.path.join(BASE_DIR, "yolo11_best.pt")
 
-        gallery.append((img, f"{fname}\n{caption}"))
+# Option B (optional): pull from a private HF model repo using a Space secret
+# Set these env vars in your Space if you want auto-download:
+#   HF_TOKEN=<read token>   YOLO_REPO_ID="yourname/yolo-detector"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+YOLO_REPO_ID = os.environ.get("YOLO_REPO_ID")
 
-        top3_labels = [ID2LABEL[i] for i in top_idxs]
-        top3_scores = [round(p[i], 4) for i in top_idxs]
-        table_rows.append([
-            fname,
-            ID2LABEL[top1],
-            round(p[top1], 4),
-            ", ".join(top3_labels),
-            ", ".join(map(str, top3_scores)),
-        ])
-
-    # Create CSV for download
-    csv_path = None
+def _download_from_hub_if_needed() -> str | None:
+    """If YOLO_REPO_ID is set, download weights with huggingface_hub; else return None."""
+    if not YOLO_REPO_ID:
+        return None
     try:
-        # Write CSV into a temp file inside project dir so Gradio can serve it
-        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="predictions_", dir=BASE_DIR, mode="w", newline='', encoding='utf-8')
-        writer = csv.writer(tmp)
-        # headers
-        writer.writerow(["filename", "top1_label", "top1_conf", "top3_labels", "top3_confs"])
-        for row in table_rows:
-            writer.writerow(row)
-        tmp.flush()
-        tmp.close()
-        csv_path = tmp.name
-    except Exception:
-        # If CSV can't be created, return None for the file but keep other outputs
-        csv_path = None
-
-    return gallery, table_rows, csv_path
-
+        from huggingface_hub import snapshot_download
+        local_dir = snapshot_download(
+            repo_id=YOLO_REPO_ID, repo_type="model", token=HF_TOKEN
+        )
+        # try common filenames
+        for name in ("yolo11_best.pt", "best.pt", "yolo.pt", "weights.pt"):
+            cand = Path(local_dir) / name
+            if cand.exists():
+                return str(cand)
+    except Exception as e:
+        print("[YOLO] Hub download failed:", e)
+    return None
 
-# ---------- NEW: YOLO object detection for multi-image upload ----------
-YOLO_WEIGHTS = os.path.join(BASE_DIR, "yolo11_best.pt")
 _yolo_model = None
 def _load_yolo():
+    """Load YOLO weights either from local file or HF Hub."""
     global _yolo_model
     if _yolo_model is not None:
         return _yolo_model
     if YOLO is None:
-        raise RuntimeError("ultralytics package not installed. Please install 'ultralytics'.")
-    if not os.path.exists(YOLO_WEIGHTS):
-        # Try current directory too
-        alt = Path.cwd() / "yolo11_best.pt"
-        if alt.exists():
-            model_path = str(alt)
-        else:
-            raise FileNotFoundError(f"YOLO weights not found at {YOLO_WEIGHTS}. Place yolo11_best.pt in project root.")
-    else:
+        raise RuntimeError("ultralytics package not installed. Add 'ultralytics' to requirements.txt")
+
+    model_path = None
+    if os.path.exists(YOLO_WEIGHTS):
         model_path = YOLO_WEIGHTS
+    else:
+        hub_path = _download_from_hub_if_needed()
+        if hub_path:
+            model_path = hub_path
+
+    if not model_path:
+        raise FileNotFoundError(
+            "YOLO weights not found. Either include 'yolo11_best.pt' in the repo root, "
+            "or set YOLO_REPO_ID (+ HF_TOKEN if private) to pull from the Hub."
+        )
 
     _yolo_model = YOLO(model_path)
     return _yolo_model
 
-
-def detect_objects_batch(files, iou=0.25, conf=0.25):
+def detect_objects_batch(files, conf=0.25, iou=0.25):
     """
-    Run YOLO detection on multiple images.
-    Returns: gallery of annotated images, dataframe rows, csv file path
+    Run YOLO detection on multiple images (up to 10).
+    Returns: gallery of annotated images, rows table, csv filepath
     """
     if YOLO is None:
         return [], [], None
-
     if not files:
         return [], [], None
 
-    # Load model
     try:
         ymodel = _load_yolo()
     except Exception as e:
         print("YOLO load error:", e)
         return [], [], None
 
-    annotated_paths = []
-    table_rows = []
-    gallery = []
+    gallery, table_rows = [], []
 
     for f in files[:MAX_BATCH]:
         path = getattr(f, "name", None) or getattr(f, "path", None) or f
         try:
-            # Run predict; returns a Results object list
             results = ymodel.predict(source=path, conf=conf, iou=iou, imgsz=640, verbose=False)
         except Exception as e:
             print(f"Detection failed for {path}:", e)
             continue
-
-        # results is list-like; take first
         res = results[0]
 
-        # Prepare annotation image using res.plot() so boxes+confidences are drawn
+        # annotated image
         ann_path = None
         try:
-            ann_img = res.plot()  # returns numpy array with annotations
-            from PIL import Image as PILImage
-            ann_pil = PILImage.fromarray(ann_img)
+            ann_img = res.plot()
+            ann_pil = Image.fromarray(ann_img)
             out_dir = tempfile.mkdtemp(prefix="yolo_out_", dir=BASE_DIR)
             os.makedirs(out_dir, exist_ok=True)
-            ann_filename = os.path.splitext(os.path.basename(path))[0] + "_annotated.jpg"
+            ann_filename = Path(path).stem + "_annotated.jpg"
             ann_path = os.path.join(out_dir, ann_filename)
             ann_pil.save(ann_path)
         except Exception:
-            # Fallback to ultralytics save if plot() isn't available
             try:
                 out_dir = tempfile.mkdtemp(prefix="yolo_out_", dir=BASE_DIR)
                 res.save(save_dir=out_dir)
-                saved_files = res.files if hasattr(res, 'files') else []
+                saved_files = getattr(res, "files", [])
                 ann_path = saved_files[0] if saved_files else None
             except Exception:
                 ann_path = None
 
-        # Build table rows from detections
-        boxes = res.boxes if hasattr(res, 'boxes') else None
+        # extract detections
+        boxes = getattr(res, "boxes", None)
         if boxes is None or len(boxes) == 0:
             table_rows.append([os.path.basename(path), 0, "", "", ""])
-            if ann_path and os.path.exists(ann_path):
-                gallery.append((Image.open(ann_path).convert('RGB'), f"{os.path.basename(path)}\nNo detections"))
-            else:
-                gallery.append((Image.open(path).convert('RGB'), f"{os.path.basename(path)}\nNo detections"))
+            img_for_gallery = Image.open(ann_path).convert("RGB") if ann_path and os.path.exists(ann_path) \
+                              else Image.open(path).convert("RGB")
+            gallery.append((img_for_gallery, f"{os.path.basename(path)}\nNo detections"))
             continue
 
-        det_labels = []
-        det_scores = []
-        det_boxes = []
+        det_labels, det_scores, det_boxes = [], [], []
         for box in boxes:
-            # box.cls, box.conf, box.xyxy
-            cls = int(box.cls.cpu().item()) if hasattr(box, 'cls') else None
-            # use .item() to extract scalar and avoid numpy deprecation warnings
-            if hasattr(box, 'conf'):
+            cls = int(box.cls.cpu().item()) if hasattr(box, "cls") else None
+            # conf
+            try:
+                confscore = float(box.conf.cpu().item()) if hasattr(box, "conf") else None
+            except Exception:
                 try:
-                    confscore = float(box.conf.cpu().item())
+                    confscore = float(box.conf.item())
                 except Exception:
-                    try:
-                        confscore = float(box.conf.item())
-                    except Exception:
-                        confscore = None
-            else:
-                confscore = None
-
-            # extract xyxy coords; box.xyxy may be shape (1,4) -> nested list after .tolist()
+                    confscore = None
+            # xyxy
             coords = []
-            if hasattr(box, 'xyxy'):
+            if hasattr(box, "xyxy"):
                 try:
                     arr = box.xyxy.cpu().numpy()
-                    # handle nested shape (1,4) or (4,)
-                    if getattr(arr, 'ndim', None) == 2 and arr.shape[0] == 1:
+                    if getattr(arr, "ndim", None) == 2 and arr.shape[0] == 1:
                         coords = arr[0].tolist()
-                    elif getattr(arr, 'ndim', None) == 1:
+                    elif getattr(arr, "ndim", None) == 1:
                         coords = arr.tolist()
                     else:
                         coords = arr.reshape(-1).tolist()
                 except Exception:
-                    # fallback: try to call tolist()
                     try:
                         coords = box.xyxy.tolist()
                     except Exception:
                         coords = []
 
-            # append detection info
             det_labels.append(ymodel.names.get(cls, str(cls)) if cls is not None else "")
             det_scores.append(round(confscore, 4) if confscore is not None else "")
-            # round and store coords
             try:
                 det_boxes.append([round(float(x), 2) for x in coords])
             except Exception:
-                # fallback: store raw repr
                 det_boxes.append([str(coords)])
 
-        # create readable label:confidence pairs
         label_conf_pairs = [f"{l}:{s}" for l, s in zip(det_labels, det_scores)]
         boxes_repr = ["[" + ", ".join(map(str, b)) + "]" for b in det_boxes]
         table_rows.append([
@@ -261,28 +167,25 @@ def detect_objects_batch(files, iou=0.25, conf=0.25):
             len(det_labels),
             ", ".join(label_conf_pairs),
             ", ".join(boxes_repr),
-            "; ".join([str(b) for b in det_boxes])
+            "; ".join([str(b) for b in det_boxes]),
         ])
 
-        # Use annotated image if exists
-        if ann_path and os.path.exists(ann_path):
-            try:
-                gallery.append((Image.open(ann_path).convert('RGB'), f"{os.path.basename(path)}\n{len(det_labels)} detections"))
-            except Exception:
-                gallery.append((Image.open(path).convert('RGB'), f"{os.path.basename(path)}\n{len(det_labels)} detections"))
-        else:
-            gallery.append((Image.open(path).convert('RGB'), f"{os.path.basename(path)}\n{len(det_labels)} detections"))
+        img_for_gallery = Image.open(ann_path).convert("RGB") if ann_path and os.path.exists(ann_path) \
+                          else Image.open(path).convert("RGB")
+        gallery.append((img_for_gallery, f"{os.path.basename(path)}\n{len(det_labels)} detections"))
 
     # write CSV
     csv_path = None
     try:
-        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="yolo_preds_", dir=BASE_DIR, mode="w", newline='', encoding='utf-8')
+        tmp = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".csv", prefix="yolo_preds_", dir=BASE_DIR,
+            mode="w", newline='', encoding='utf-8'
+        )
         writer = csv.writer(tmp)
         writer.writerow(["filename", "num_detections", "labels_with_conf", "boxes", "raw_boxes"])
         for r in table_rows:
             writer.writerow(r)
-        tmp.flush()
-        tmp.close()
+        tmp.flush(); tmp.close()
         csv_path = tmp.name
     except Exception as e:
         print("Failed to write CSV:", e)
@@ -291,49 +194,35 @@ def detect_objects_batch(files, iou=0.25, conf=0.25):
     return gallery, table_rows, csv_path
 
 # ---------- UI ----------
-single = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload Underwater Image"),
-    outputs=gr.Label(num_top_classes=len(ID2LABEL), label="Species Classification"),
-    title="🌊 BenthicAI - Single Image",
-    description="Classify one image into one of 7 benthic species."
-)
-
-batch = gr.Interface(
-    fn=classify_images_batch,
-    inputs=gr.Files(label="Upload up to 10 images"),
-    outputs=[
-        gr.Gallery(label="Results (Top-1 in caption)", height=500, rows=3),
-        gr.Dataframe(
-            headers=["filename", "top1_label", "top1_conf", "top3_labels", "top3_confs"],
-            label="Predictions Table",
-            wrap=True
-        )
-        , gr.File(label="Download CSV")
-    ],
-    title="🌊 BenthicAI - Batch (up to 10)",
-    description="Upload multiple images (max 10). Outputs a gallery with captions and a table of top predictions.",
-)
-
-demo = gr.TabbedInterface([single, batch], ["Single", "Batch"])
-print(YOLO==None, flush=True)
-# Add Object Detection tab if ultralytics available
-if YOLO is not None:
-    detection_iface = gr.Interface(
+if YOLO is None:
+    demo = gr.Interface(
+        fn=lambda *a, **k: ("Ultralytics not installed; add 'ultralytics' to requirements.txt",),
+        inputs=[],
+        outputs="text",
+        title="🌊 BenthicAI — Object Detection",
+        description="Ultralytics is not installed."
+    )
+else:
+    demo = gr.Interface(
         fn=detect_objects_batch,
-        inputs=[gr.Files(label="Upload images for detection (max 10)"), gr.Slider(minimum=0.0, maximum=1.0, value=0.25, label="conf threshold"), gr.Slider(minimum=0.0, maximum=1.0, value=0.25, label="IOU threshold")],
+        inputs=[
+            gr.Files(label="Upload images (max 10)"),
+            gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.01, label="Conf threshold"),
+            gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.01, label="IoU threshold"),
+        ],
         outputs=[
             gr.Gallery(label="Detections (annotated)", height=500, rows=3),
-            gr.Dataframe(headers=["filename", "num_detections", "labels_with_conf", "boxes", "raw_boxes"], label="Detection Table"),
-            gr.File(label="Download CSV")
+            gr.Dataframe(headers=["filename", "num_detections", "labels_with_conf", "boxes", "raw_boxes"],
+                         label="Detection Table"),
+            gr.File(label="Download CSV"),
         ],
-        title="🌊 BenthicAI - Object Detection",
-        description="Run YOLO object detection on multiple images. Requires 'yolo11_best.pt' in project root."
+        title="🌊 BenthicAI — Object Detection",
+        description=(
+            "Run YOLO object detection on multiple images. "
+            "Place 'yolo11_best.pt' in the repo root, OR set YOLO_REPO_ID (+ HF_TOKEN if private) "
+            "to fetch from the Hub."
+        ),
     )
 
-    # extend tabs
-    demo = gr.TabbedInterface([single, batch, detection_iface], ["Single", "Batch", "Detection"])
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
-
+    demo.launch(server_name="0.0.0.0", server_port=7860)