add object detection

app.py

@@ -1,115 +1,4 @@
-# import gradio as gr
-# import torch
-# from transformers import AutoImageProcessor, AutoModelForImageClassification
-# from PIL import Image
-# import os
-
-# # Get model path (Windows compatible)
-# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# MODEL_ID = "dshi01/convnext-tiny-224-7clss"
-
-# # Try different possible filenames
-# # possible_names = ['ConvNextmodel.pth', 'convnextmodel.pth', 'ConvNext_model.pth']
-# # model_path = None
-
-# # for name in possible_names:
-# #     test_path = os.path.join(BASE_DIR, name)
-# #     if os.path.exists(test_path):
-# #         model_path = test_path
-# #         print(f"✓ Found model: {name}")
-# #         break
-
-# # if model_path is None:
-# #     raise FileNotFoundError(f"Could not find model file. Tried: {possible_names}")
-
-# # Species categories (7 classes)
-# SPECIES_CATEGORIES = [
-#     'Eel',
-#     'Scallop',
-#     'Crab',
-#     'Flatfish',
-#     'Roundfish',
-#     'Skate',
-#     'Whelk'
-# ]
-
-# # Load model
-# print(f"Loading model from: {MODEL_ID}")
-# # model = AutoModelForImageClassification.from_pretrained(
-# #     'facebook/convnext-tiny-224',
-# #     num_labels=7,
-# #     ignore_mismatched_sizes=True
-# # )
-# processor=AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
-# model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
-
-# # Load weights
-# # checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
-# # if isinstance(checkpoint, dict):
-# #     if 'model' in checkpoint:
-# #         checkpoint = checkpoint['model']
-# #     elif 'state_dict' in checkpoint:
-# #         checkpoint = checkpoint['state_dict']
-
-# # model.load_state_dict(checkpoint, strict=False)
-# # model.eval()
-
-# # Load processor
-# # processor = AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
-# # print("✓ Model loaded successfully!")
-
-# def classify_image(image):
-#     """
-#     Classify a benthic species image.
-    
-#     Args:
-#         image: PIL Image or numpy array
-    
-#     Returns:
-#         dict: Predictions with species names and confidence scores
-#     """
-#     # Convert to PIL if needed
-#     if not isinstance(image, Image.Image):
-#         image = Image.fromarray(image).convert('RGB')
-    
-#     # Preprocess
-#     inputs = processor(images=image, return_tensors="pt")
-    
-#     # Predict
-#     with torch.no_grad():
-#         outputs = model(**inputs)
-#         logits = outputs.logits
-#         probabilities = torch.nn.functional.softmax(logits, dim=1)
-    
-#     # Create results dictionary for Gradio
-#     results = {}
-#     for idx, prob in enumerate(probabilities[0]):
-#         results[SPECIES_CATEGORIES[idx]] = float(prob)
-    
-#     return results
-
-# # Create Gradio interface
-# demo = gr.Interface(
-#     fn=classify_image,
-#     inputs=gr.Image(type="pil", label="Upload Underwater Image"),
-#     outputs=gr.Label(num_top_classes=7, label="Species Classification"),
-#     title="🌊 BenthicAI - Benthic Species Classifier",
-#     description="Upload an image of a benthic organism to classify it into one of 7 species categories. Built with ConvNeXT transformer model.",
-#     examples=[
-#         [os.path.join("examples", "eel.jpg")],
-#         [os.path.join("examples", "scallop.jpg")],
-#         [os.path.join("examples", "crab.jpg")],
-#     ] if os.path.exists("examples") else None,
-#     theme=gr.themes.Soft(),
-#     allow_flagging="never"
-# )
-
-# if __name__ == "__main__":
-#     demo.launch(
-#         server_name="0.0.0.0",
-#         server_port=7860,
-#         share=True  # Set to True to get a public URL
-#     )
+#Main Gradio app ith image classification and object detection tabs
 import gradio as gr
 import torch
 import torch.nn.functional as F
@@ -118,6 +7,13 @@ 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"
@@ -221,6 +117,179 @@ def classify_images_batch(files):
 
     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,
@@ -247,6 +316,23 @@ batch = gr.Interface(
 )
 
 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

@@ -2,4 +2,5 @@ torch
 torchvision
 transformers
 gradio
-Pillow
+Pillow
+ultralytics