Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image | |
| import tempfile | |
| import os | |
| import subprocess | |
| import glob | |
| def process_image(input_img): | |
| # Create a temporary directory to store the input image | |
| with tempfile.TemporaryDirectory() as temp_input_dir: | |
| input_image_path = os.path.join(temp_input_dir, "input.jpg") | |
| input_img.save(input_image_path) | |
| # Create a temporary directory for the output image | |
| with tempfile.TemporaryDirectory() as temp_output_dir: | |
| # Command to run the YOLO model | |
| command = f"yolo task=detect mode=predict model=best.pt conf=0.25 source={temp_input_dir} save=True" | |
| subprocess.run(command, shell=True) | |
| # Get the most recent 'predict' folder in 'runs/detect' | |
| list_of_dirs = glob.glob('runs/detect/predict*') | |
| latest_dir = max(list_of_dirs, key=os.path.getctime) | |
| # Assuming YOLO saves the output with the same name in the latest 'predict' folder | |
| output_image_name = os.path.basename(input_image_path) | |
| output_image_path = os.path.join(latest_dir, output_image_name) | |
| if os.path.exists(output_image_path): | |
| output_img = Image.open(output_image_path) | |
| return output_img | |
| else: | |
| return "No output image found." | |
| # Define the Gradio interface | |
| demo = gr.Interface( | |
| fn=process_image, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Image(type="pil"), | |
| title="Object Detection with YOLO", | |
| description="Upload an image and the YOLO model will detect objects." | |
| ) | |
| # Launch the app | |
| demo.launch(share=True) | |