File size: 3,960 Bytes
dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 dff62de c991e43 e5f266f c317c06 dff62de c317c06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import gradio as gr
import torch
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)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
weights = EfficientNet_V2_S_Weights.IMAGENET1K_V1
model = efficientnet_v2_s(weights=weights)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval().to(device)
mean = getattr(weights, "meta", {}).get("mean", [0.485, 0.456, 0.406])
std = getattr(weights, "meta", {}).get("std", [0.229, 0.224, 0.225])
transform = transforms.Compose([
transforms.Resize((384, 384)),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
])
# ---------------------------
# 2. Load YOLOv8 model
# ---------------------------
yolo_model = YOLO("yolov8x.pt")
# ---------------------------
# 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)
# Run YOLOv8 detection
results = yolo_model(image_cv)
boxes = results[0].boxes
names = results[0].names
detections = boxes.data.cpu().numpy()
# 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)
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=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 - 12/2/2025 Build",
description="Upload an image with one or more cats.",
api_name="predict"
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0")
|