Spaces:
Sleeping
Sleeping
File size: 3,096 Bytes
69fe549 |
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 |
import torch
import gradio as gr
import numpy as np
import pandas as pd
from ultralytics import YOLO
import cv2
# Load the trained YOLO model
model = YOLO('best.pt')
def predict_image(image):
"""
Perform object detection on the input image
Args:
image (numpy.ndarray): Input image
Returns:
tuple: Annotated image, results dataframe, precision-recall table
"""
# Perform prediction
results = model(image)[0]
# Create a copy of the image for annotation
annotated_image = image.copy()
# Prepare results for display
detection_results = []
# Draw bounding boxes and collect detection info
for box in results.boxes:
# Extract box coordinates
x1, y1, x2, y2 = box.xyxy[0].int().tolist()
# Get class and confidence
cls = int(box.cls[0])
conf = float(box.conf[0])
class_name = model.names[cls]
# Draw bounding box
cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
# Put text on the image
label = f'{class_name} {conf:.2f}'
cv2.putText(annotated_image, label, (x1, y1-10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Collect detection details
detection_results.append({
'Class': class_name,
'Confidence': f'{conf:.4f}',
'Bounding Box': f'[{x1},{y1},{x2},{y2}]'
})
# Create DataFrame for detections
detections_df = pd.DataFrame(detection_results)
# Generate precision-recall metrics (placeholder - you'll need to replace with actual metrics)
pr_metrics = generate_pr_metrics()
return annotated_image, detections_df, pr_metrics
def generate_pr_metrics():
"""
Generate placeholder precision and recall metrics
Note: Replace this with actual calculations from your training results
"""
# This is a placeholder - you should replace with actual metrics from your training
metrics_data = {
'Class': model.names.values(),
'Precision': [0.85, 0.90, 0.88], # Example values
'Recall': [0.80, 0.85, 0.82] # Example values
}
metrics_df = pd.DataFrame(metrics_data)
return metrics_df
# Create Gradio interface
def launch_interface():
# Input image upload
image_input = gr.Image(type="numpy", label="Upload Image")
# Outputs
annotated_image_output = gr.Image(label="Annotated Image")
detections_output = gr.Dataframe(label="Detection Results")
metrics_output = gr.Dataframe(label="Precision-Recall Metrics")
# Gradio interface
demo = gr.Interface(
fn=predict_image,
inputs=image_input,
outputs=[
annotated_image_output,
detections_output,
metrics_output
],
title="YOLO Object Detection",
description="Upload an image for object detection"
)
return demo
# Launch the app
if __name__ == "__main__":
demo = launch_interface()
demo.launch() |