Added yolov8, updated app.py

app.py

@@ -4,9 +4,15 @@ import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights
 from PIL import Image
+import numpy as np
+import cv2
+from ultralytics import YOLO
 import json
 import os
 
+# ---------------------------
+# 1. Load EfficientNet model
+# ---------------------------
 MODEL_PATH = "best_model.pth"
 with open("class_names.json", "r") as f:
     CLASS_NAMES = json.load(f)
@@ -28,29 +34,94 @@ transform = transforms.Compose([
     transforms.Normalize(mean=mean, std=std),
 ])
 
-def predict(image):
-    image = transform(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        outputs = model(image)
-        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
+# ---------------------------
+# 2. Load YOLOv8 model
+# ---------------------------
+yolo_model = YOLO("yolov8x.pt")
 
-    results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
+# ---------------------------
+# 3. Detection + Identification Pipeline
+# ---------------------------
+def detect_and_identify(image: Image.Image):
+    # Convert PIL to OpenCV (numpy)
+    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-    sorted_indices = torch.argsort(probs, descending=True)
-    top5_indices = sorted_indices[:5]
+    # Run YOLOv8 detection
+    results = yolo_model(image_cv)
+    boxes = results[0].boxes
+    names = results[0].names
+    detections = boxes.data.cpu().numpy()
 
-    top5_results = {CLASS_NAMES[i]: float(probs[i]) for i in top5_indices}
+    # Filter for cats only
+    valid_detections = []
+    for det in detections:
+        x1, y1, x2, y2, conf, cls = det
+        class_name = names[int(cls)]
+        if class_name == "cat":
+            valid_detections.append(det)
 
-    return top5_results, results
+    if not valid_detections:
+        return (
+            None,
+            {"Error": "No cat detected in the image."},
+            {"Info": "Please upload an image containing one or more visible cats."}
+        )
 
+    cropped_images = []
+    predictions = {}
 
+    for i, det in enumerate(valid_detections):
+        x1, y1, x2, y2, conf, cls = det
+        x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+
+        # Crop the detected cat
+        cropped_cat = image_cv[y1:y2, x1:x2]
+        cropped_pil = Image.fromarray(cv2.cvtColor(cropped_cat, cv2.COLOR_BGR2RGB))
+
+        # Preprocess for classifier
+        input_tensor = transform(cropped_pil).unsqueeze(0).to(device)
+
+        # Identify using EfficientNet
+        with torch.no_grad():
+            outputs = model(input_tensor)
+            probs = torch.nn.functional.softmax(outputs, dim=1)[0]
+
+        # Get top prediction
+        top_idx = torch.argmax(probs).item()
+        top_label = CLASS_NAMES[top_idx]
+        confidence = float(probs[top_idx])
+
+        # Get Top 5 predictions only
+        top5_indices = torch.argsort(probs, descending=True)[:5]
+        top5_results = {CLASS_NAMES[j]: float(probs[j]) for j in top5_indices}
+
+        # Append image with caption (for Gallery)
+        cropped_images.append((cropped_pil, f"{top_label} ({confidence*100:.1f}%)"))
+
+        # Add JSON data for this cat
+        predictions[f"Cat {i+1} - {top_label}"] = {
+            "Top Prediction": top_label,
+            "Confidence": confidence,
+            "Top 5 Scores": top5_results
+        }
+
+
+    return cropped_images, predictions
+
+
+# ---------------------------
+# 4. Gradio Interface
+# ---------------------------
 demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Label(label="Prediction"), gr.JSON(label="Confidence Scores")],
+    fn=detect_and_identify,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[
+        gr.Gallery(label="Detected Cats (Cropped)", columns=2, rows=2),
+        gr.JSON(label="Predictions per Cat")
+    ],
     title="StrAI - Cat Identifier",
-    description="Upload an image to identify which cat it is."
+    description="Upload an image with one or more cats."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

yolov8x.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3df4ada6b4dad6d657868f2fdf7faecfb34dcfccf3a25c4b82079064718524c8
+size 136890692