File size: 3,055 Bytes
d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b d2ff043 110c98b | 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 | import gradio as gr
import requests
import numpy as np
from PIL import Image
import io
from ultralytics import YOLO
import cv2
# Load the YOLO model at startup
try:
model = YOLO('modelo_epoch_50.pt')
print("Model loaded successfully")
except Exception as e:
print(f"Error loading model: {str(e)}")
model = None
# Function to process the captured image
def process_image(image):
# Check if image is None
if image is None:
raise gr.Error("Please take a picture first before analyzing!")
# Convert image to RGB if it's not already
if image.mode != 'RGB':
image = image.convert('RGB')
# Run inference
results = model.predict(image, save=True, conf=0.5)
# Print the results
print("Model predictions:", results[0].boxes)
# Convert PIL Image to numpy array for OpenCV
image_cv = np.array(image)
# Convert RGB to BGR (OpenCV uses BGR format)
image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
# Draw bounding boxes of the prediction
boxes = results[0].boxes
for box in boxes:
b = box.xyxy[0] # Bounding box coordinates
c = box.cls # Predicted class
confidence = box.conf
x1, y1, x2, y2 = map(int, b)
cv2.rectangle(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2) # Blue for prediction
label = f"{results[0].names[int(c)]} {confidence.item():.2f}"
cv2.putText(image_cv, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
print(label)
# Convert back to RGB for display if needed
image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
# Convert back to PIL Image
final_image = Image.fromarray(image_cv)
# Save the final image to current directory, overwriting if exists
final_image.save('FINAL.jpg', 'JPEG', quality=95)
return final_image
# Create the Gradio interface
with gr.Blocks() as demo:
with gr.Row():
# Left column for camera and controls
with gr.Column(scale=1):
gr.Markdown("### Step 1: Take a Picture")
camera = gr.Image(type="pil", label="Camera View", sources=["webcam"])
gr.Markdown("Click the '📸 Take Photo' button in the camera view above")
with gr.Row():
analyze_btn = gr.Button("🔍 Analyze", variant="primary")
new_picture_btn = gr.Button("🔄 Reset")
# Right column for results
with gr.Column(scale=1):
gr.Markdown("### Step 2: View Results")
modified_image = gr.Image(label="Analyzed Image")
# Set up the event handlers
analyze_btn.click(
fn=process_image,
inputs=[camera],
outputs=[modified_image]
)
new_picture_btn.click(
fn=lambda: (None, None),
inputs=[],
outputs=[camera, modified_image]
)
# Launch the interface
if __name__ == "__main__":
demo.launch(share=True) |