Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import json | |
| import os | |
| from pathlib import Path | |
| from PIL import Image | |
| import shutil | |
| from ultralytics import YOLO | |
| def load_models(models_dir='models', info_file='models_info.json'): | |
| """ | |
| Load YOLO models and their information from the specified directory and JSON file. | |
| Args: | |
| models_dir (str): Path to the models directory. | |
| info_file (str): Path to the JSON file containing model info. | |
| Returns: | |
| dict: A dictionary of models and their associated information. | |
| """ | |
| with open(info_file, 'r') as f: | |
| models_info = json.load(f) | |
| models = {} | |
| for model_info in models_info: | |
| model_name = model_info['model_name'] | |
| model_path = os.path.join(models_dir, model_name, 'best.pt') # Assuming 'best.pt' as the weight file | |
| if os.path.isfile(model_path): | |
| try: | |
| # Load the YOLO model | |
| model = YOLO(model_path) | |
| models[model_name] = { | |
| 'model': model, | |
| 'mAP': model_info.get('mAP_score', 'N/A'), | |
| 'num_images': model_info.get('num_images', 'N/A') | |
| } | |
| print(f"Loaded model '{model_name}' from '{model_path}'.") | |
| except Exception as e: | |
| print(f"Error loading model '{model_name}': {e}") | |
| else: | |
| print(f"Model weight file for '{model_name}' not found at '{model_path}'. Skipping.") | |
| return models | |
| def get_model_info(model_name, models): | |
| """ | |
| Retrieve model information for the selected model. | |
| Args: | |
| model_name (str): The name of the model. | |
| models (dict): The dictionary containing models and their info. | |
| Returns: | |
| str: A formatted string containing model information. | |
| """ | |
| model_info = models.get(model_name, {}) | |
| if not model_info: | |
| return "Model information not available." | |
| info_text = ( | |
| f"**Model Name:** {model_name}\n\n" | |
| f"**mAP Score:** {model_info.get('mAP', 'N/A')}\n\n" | |
| f"**Number of Images Trained On:** {model_info.get('num_images', 'N/A')}" | |
| ) | |
| return info_text | |
| def predict_image(model_name, image, models): | |
| """ | |
| Perform prediction on an uploaded image using the selected YOLO model. | |
| Args: | |
| model_name (str): The name of the selected model. | |
| image (PIL.Image.Image): The uploaded image. | |
| models (dict): The dictionary containing models and their info. | |
| Returns: | |
| tuple: A status message, the processed image, and the path to the output image. | |
| """ | |
| model = models.get(model_name, {}).get('model', None) | |
| if not model: | |
| return "Error: Model not found.", None, None | |
| try: | |
| # Save the uploaded image to a temporary path | |
| input_image_path = f"temp/{model_name}_input_image.jpg" | |
| os.makedirs(os.path.dirname(input_image_path), exist_ok=True) | |
| image.save(input_image_path) | |
| # Perform prediction | |
| results = model(input_image_path, save=True, save_txt=False, conf=0.25) | |
| # Ultralytics saves the result images in 'runs/detect/predict' | |
| output_image_path = results[0].save()[0] # Get the path to the saved image | |
| # Open the output image | |
| output_image = Image.open(output_image_path) | |
| return "Prediction completed successfully.", output_image, output_image_path | |
| except Exception as e: | |
| return f"Error during prediction: {str(e)}", None, None | |
| def predict_video(model_name, video, models): | |
| """ | |
| Perform prediction on an uploaded video using the selected YOLO model. | |
| Args: | |
| model_name (str): The name of the selected model. | |
| video (str): Path to the uploaded video file. | |
| models (dict): The dictionary containing models and their info. | |
| Returns: | |
| tuple: A status message, the processed video, and the path to the output video. | |
| """ | |
| model = models.get(model_name, {}).get('model', None) | |
| if not model: | |
| return "Error: Model not found.", None, None | |
| try: | |
| # Ensure the video is saved in a temporary location | |
| input_video_path = video.name | |
| if not os.path.isfile(input_video_path): | |
| # If the video is a temp file provided by Gradio | |
| shutil.copy(video.name, input_video_path) | |
| # Perform prediction | |
| results = model(input_video_path, save=True, save_txt=False, conf=0.25) | |
| # Ultralytics saves the result videos in 'runs/detect/predict' | |
| output_video_path = results[0].save()[0] # Get the path to the saved video | |
| return "Prediction completed successfully.", output_video_path, output_video_path | |
| except Exception as e: | |
| return f"Error during prediction: {str(e)}", None, None | |
| def main(): | |
| # Load the models and their information | |
| models = load_models() | |
| # Initialize Gradio Blocks interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🧪 YOLO Model Tester") | |
| gr.Markdown( | |
| """ | |
| Upload images or videos to test different YOLO models. Select a model from the dropdown to see its details. | |
| """ | |
| ) | |
| # Model selection and info | |
| with gr.Row(): | |
| model_dropdown = gr.Dropdown( | |
| choices=list(models.keys()), | |
| label="Select Model", | |
| value=None | |
| ) | |
| model_info = gr.Markdown("**Model Information will appear here.**") | |
| # Update model_info when a model is selected | |
| model_dropdown.change( | |
| fn=lambda model_name: get_model_info(model_name, models) if model_name else "Please select a model.", | |
| inputs=model_dropdown, | |
| outputs=model_info | |
| ) | |
| # Tabs for different input types | |
| with gr.Tabs(): | |
| # Image Prediction Tab | |
| with gr.Tab("🖼️ Image"): | |
| with gr.Column(): | |
| image_input = gr.Image( | |
| type='pil', | |
| label="Upload Image for Prediction", | |
| tool="editor" | |
| ) | |
| image_predict_btn = gr.Button("🔍 Predict on Image") | |
| image_status = gr.Markdown("**Status will appear here.**") | |
| image_output = gr.Image(label="Predicted Image") | |
| image_download_btn = gr.File(label="⬇️ Download Predicted Image") | |
| # Define the image prediction function | |
| def process_image(model_name, image): | |
| return predict_image(model_name, image, models) | |
| # Connect the predict button | |
| image_predict_btn.click( | |
| fn=process_image, | |
| inputs=[model_dropdown, image_input], | |
| outputs=[image_status, image_output, image_download_btn] | |
| ) | |
| # Video Prediction Tab | |
| with gr.Tab("🎥 Video"): | |
| with gr.Column(): | |
| video_input = gr.Video( | |
| label="Upload Video for Prediction" | |
| ) | |
| video_predict_btn = gr.Button("🔍 Predict on Video") | |
| video_status = gr.Markdown("**Status will appear here.**") | |
| video_output = gr.Video(label="Predicted Video") | |
| video_download_btn = gr.File(label="⬇️ Download Predicted Video") | |
| # Define the video prediction function | |
| def process_video(model_name, video): | |
| return predict_video(model_name, video, models) | |
| # Connect the predict button | |
| video_predict_btn.click( | |
| fn=process_video, | |
| inputs=[model_dropdown, video_input], | |
| outputs=[video_status, video_output, video_download_btn] | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| **Note:** Ensure that the YOLO models are correctly placed in the `models/` directory and that `models_info.json` is properly configured. | |
| """ | |
| ) | |
| # Launch the Gradio app | |
| demo.launch() | |
| if __name__ == "__main__": | |
| main() | |