Spaces:
Sleeping
Sleeping
File size: 1,690 Bytes
80d1a24 |
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 |
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)
|