File size: 2,722 Bytes
d9a911c ffd04a0 d9a911c 3a7a4cb d9a911c 3a7a4cb d9a911c 0c26def d9a911c 3a7a4cb d9a911c 3a7a4cb d9a911c 3a7a4cb d9a911c 3a7a4cb d9a911c 3a7a4cb d9a911c 0c26def | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import gradio as gr
from ultralytics import YOLO
import os
import shutil
# 1. Load the fine-tuned YOLO model
model_path = "best.pt"
model = YOLO(model_path)
# Define the directory to save inference results
output_dir = "inference_results"
os.makedirs(output_dir, exist_ok=True)
# 2. Define an inference function for images
def predict_image(image_file, conf_threshold):
if image_file is None:
return "No input image provided."
print(f"Processing: {image_file}")
print(f"Confidence threshold: {conf_threshold}")
# Clear previous results if any
for item in os.listdir(output_dir):
item_path = os.path.join(output_dir, item)
if os.path.isfile(item_path):
os.remove(item_path)
elif os.path.isdir(item_path):
shutil.rmtree(item_path)
results = model.predict(source=image_file, conf=conf_threshold, save=True, project=output_dir, name="run", exist_ok=True)
# The results are saved in a subdirectory within output_dir/run/
# We need to find the actual path to the saved file.
# Ultralytics saves to runs/detect/runX (where X is an incrementing number)
# We need to find the latest run folder.
# Get the latest run folder within the output_dir
run_folders = [d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))]
run_folders.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
if not run_folders:
return "No detection results saved."
latest_run_path = os.path.join(output_dir, run_folders[0])
# Now find the actual image file inside this run_folder
detected_files = [f for f in os.listdir(latest_run_path) if f.endswith(('.jpg', '.jpeg', '.png'))]
if not detected_files:
return "No detected image found in results."
# Assuming only one file is processed at a time
result_path = os.path.join(latest_run_path, detected_files[0])
print(f"Results saved to: {result_path}")
return result_path
# 3. Create a Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# YOLOv8 Signature Detection")
gr.Markdown("Upload an image to perform signature detection using a fine-tuned YOLOv8n model.")
with gr.Tab("Image Detection"):
image_input = gr.Image(type="filepath", label="Upload Image")
image_conf_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold")
image_output = gr.Image(label="Detection Results")
image_button = gr.Button("Detect Signature in Image")
image_button.click(predict_image, inputs=[image_input, image_conf_slider], outputs=image_output)
# 4. Launch the Gradio interface
demo.launch(share=True) |