import os import zipfile from ultralytics import YOLO from PIL import Image import gradio as gr # Unzip the model file if it's not already unzipped zip_file = "yolov8_tuned_with_weightsv4.zip" model_file = "yolov8_tuned_with_weightsv4.pt" if not os.path.exists(model_file): with zipfile.ZipFile(zip_file, "r") as zip_ref: zip_ref.extractall() # Extract to current directory # Load the YOLOv8 model model = YOLO(model_file) # Define the class names class_names = {0: "Dumbbell", 1: "Kettlebell"} # Adjust based on your dataset # Define prediction function def predict(image): # Perform inference results = model.predict(image) output_image = results[0].plot() # Annotated image with predictions # Prepare detection details detections = [ { "class": class_names.get(int(result.cls), "Unknown"), # Map class index to name "confidence": float(result.conf), "bbox": result.xyxy.tolist(), } for result in results[0].boxes ] return output_image, detections # Create Gradio interface interface = gr.Interface( fn=predict, inputs=gr.Image(type="pil", label="Upload an Image"), # File upload button outputs=[ gr.Image(type="pil", label="Predicted Image"), # Display annotated image gr.JSON(label="Detection Details"), # Show detection results as JSON ], title="YOLOv8 Dumbbell and Kettlebell Detection", description="Upload an image to detect dumbbells and kettlebells using YOLOv8.", ) # Launch the app if __name__ == "__main__": interface.launch()