| 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 |
|
|
| |
| |
| |
| 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), |
| ]) |
|
|
| |
| |
| |
| yolo_model = YOLO("yolov8x.pt") |
|
|
| |
| |
| |
| def detect_and_identify(image: Image.Image): |
| |
| image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) |
|
|
| |
| results = yolo_model(image_cv) |
| boxes = results[0].boxes |
| names = results[0].names |
| detections = boxes.data.cpu().numpy() |
|
|
| |
| 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]) |
|
|
| |
| cropped_cat = image_cv[y1:y2, x1:x2] |
| cropped_pil = Image.fromarray(cv2.cvtColor(cropped_cat, cv2.COLOR_BGR2RGB)) |
|
|
| |
| input_tensor = transform(cropped_pil).unsqueeze(0).to(device) |
|
|
| |
| with torch.no_grad(): |
| outputs = model(input_tensor) |
| probs = torch.nn.functional.softmax(outputs, dim=1)[0] |
|
|
| |
| top_idx = torch.argmax(probs).item() |
| top_label = CLASS_NAMES[top_idx] |
| confidence = float(probs[top_idx]) |
|
|
| |
| top5_indices = torch.argsort(probs, descending=True)[:5] |
| top5_results = {CLASS_NAMES[j]: float(probs[j]) for j in top5_indices} |
|
|
| |
| cropped_images.append((cropped_pil, f"{top_label} ({confidence*100:.1f}%)")) |
|
|
| |
| predictions[f"Cat {i+1} - {top_label}"] = { |
| "Top Prediction": top_label, |
| "Confidence": confidence, |
| "Top 5 Scores": top5_results |
| } |
|
|
|
|
| return cropped_images, predictions |
|
|
|
|
| |
| |
| |
| 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") |
|
|
