obj detect app

app.py

@@ -0,0 +1,339 @@
+#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
+import os
+import csv
+import tempfile
+from pathlib import Path
+from ultralytics import YOLO
+# ultralytics YOLO import (for object detection)
+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%})"
+
+        gallery.append((img, f"{fname}\n{caption}"))
+
+        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
+    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
+
+
+# ---------- NEW: YOLO object detection for multi-image upload ----------
+YOLO_WEIGHTS = os.path.join(BASE_DIR, "yolo11_best.pt")
+_yolo_model = None
+def _load_yolo():
+    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:
+        model_path = YOLO_WEIGHTS
+
+    _yolo_model = YOLO(model_path)
+    return _yolo_model
+
+
+def detect_objects_batch(files, iou=0.25, conf=0.25):
+    """
+    Run YOLO detection on multiple images.
+    Returns: gallery of annotated images, dataframe rows, csv file path
+    """
+    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 = []
+
+    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
+        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)
+            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_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 []
+                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
+        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"))
+            continue
+
+        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'):
+                try:
+                    confscore = float(box.conf.cpu().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()
+            coords = []
+            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:
+                        coords = arr[0].tolist()
+                    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([
+            os.path.basename(path),
+            len(det_labels),
+            ", ".join(label_conf_pairs),
+            ", ".join(boxes_repr),
+            "; ".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"))
+
+    # 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')
+        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()
+        csv_path = tmp.name
+    except Exception as e:
+        print("Failed to write CSV:", e)
+        csv_path = None
+
+    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(
+        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")],
+        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")
+        ],
+        title="🌊 BenthicAI - Object Detection",
+        description="Run YOLO object detection on multiple images. Requires 'yolo11_best.pt' in project root."
+    )
+
+    # 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)
+

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+transformers
+gradio
+Pillow
+ultralytics