Spaces:
Sleeping
Sleeping
File size: 3,106 Bytes
ca0c0af |
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 |
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import os
from PIL import Image
# Load model
model_path = "models/aov_herodetector_v1.pt"
try:
model = YOLO(model_path)
print(f"Model loaded successfully from {model_path}")
except Exception as e:
print(f"Error loading model: {e}")
model = None
def detect_heroes(image):
"""
Detect AOV heroes in the uploaded image
"""
if model is None:
return None, "Model not loaded properly"
try:
# Convert PIL Image to numpy array if needed
if isinstance(image, Image.Image):
image = np.array(image)
# Run inference
results = model(image)
# Get the first result
result = results[0]
# Plot results on image
annotated_image = result.plot()
# Convert BGR to RGB (OpenCV uses BGR, but gradio expects RGB)
annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
# Get detection info
detection_info = ""
if len(result.boxes) > 0:
detection_info = f"Detected {len(result.boxes)} heroes:\n"
for i, box in enumerate(result.boxes):
conf = box.conf[0].item()
cls = int(box.cls[0].item())
class_name = model.names[cls] if cls < len(model.names) else f"Class_{cls}"
detection_info += f"- {class_name}: {conf:.2f}\n"
else:
detection_info = "No heroes detected"
return annotated_image, detection_info
except Exception as e:
return None, f"Error during detection: {str(e)}"
# Create Gradio interface
with gr.Blocks(title="AOV Hero Detector", theme=gr.themes.Soft()) as demo:
gr.Markdown("# 🎮 Arena of Valor Hero Detector")
gr.Markdown("Upload an image to detect AOV heroes using YOLOv8 model")
with gr.Row():
with gr.Column():
input_image = gr.Image(
label="Upload Image",
type="pil",
height=400
)
detect_btn = gr.Button("🔍 Detect Heroes", variant="primary")
with gr.Column():
output_image = gr.Image(
label="Detection Result",
height=400
)
detection_text = gr.Textbox(
label="Detection Info",
lines=5,
max_lines=10
)
# Examples (optional - you can add sample images)
gr.Examples(
examples=[], # Add paths to example images if you have them
inputs=input_image,
label="Example Images"
)
# Event handlers
detect_btn.click(
fn=detect_heroes,
inputs=input_image,
outputs=[output_image, detection_text]
)
# Auto-detect when image is uploaded
input_image.change(
fn=detect_heroes,
inputs=input_image,
outputs=[output_image, detection_text]
)
# Launch the app
if __name__ == "__main__":
demo.launch() |