Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import cv2 | |
| # Define the paths to the weights files | |
| MODEL_PATHS = { | |
| "YOLO - 100 Epochs": "100best.pt", | |
| "YOLO - 150 Epochs": "150best.pt", | |
| "YOLO - 200 Epochs": "200best.pt" | |
| } | |
| def detect_and_count(image, model_choice): | |
| # Dynamically load the weights for the selected choice | |
| weights_path = MODEL_PATHS[model_choice] | |
| selected_model = YOLO(weights_path) | |
| # Run prediction | |
| results = selected_model.predict(source=image, conf=0.30, iou=0.45) | |
| result = results[0] | |
| # Extract structural analytics | |
| object_count = len(result.boxes) | |
| annotated_image = result.plot() | |
| summary_text = f"Active Weights: {weights_path}\nTotal objects detected: {object_count}" | |
| return annotated_image, summary_text | |
| # Define the user interface | |
| interface = gr.Interface( | |
| fn=detect_and_count, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload an Image"), | |
| gr.Dropdown( | |
| choices=["YOLO - 100 Epochs", "YOLO - 150 Epochs", "YOLO - 200 Epochs"], | |
| value="YOLO - 100 Epochs", | |
| label="Select Training Duration (Epochs)" | |
| ) | |
| ], | |
| outputs=[ | |
| gr.Image(type="numpy", label="Detection Output"), | |
| gr.Textbox(label="Performance Summary") | |
| ], | |
| title="Multi-Epoch Comparison Object Detector", | |
| description="Upload an image and switch between the 100, 150, and 200 epoch models. Click 'Submit' after switching to update." | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() |