|
import gradio as gr |
|
import os |
|
import subprocess |
|
import cv2 |
|
import numpy as np |
|
|
|
|
|
def detect_and_crop(input_image): |
|
|
|
weights_path = 'yolo/yolov7-main/best.pt' |
|
img_size = 640 |
|
conf = 0.20 |
|
source = 'dataset/images/train/' |
|
|
|
|
|
os.makedirs(source, exist_ok=True) |
|
|
|
|
|
input_image.save(os.path.join(source, 'input_image.jpg')) |
|
|
|
|
|
command = [ |
|
'python', 'yolo/yolov7-main/detect.py', |
|
'--weights', weights_path, |
|
'--conf-thres', str(conf), |
|
'--img-size', str(img_size), |
|
'--source', os.path.join(source, 'input_image.jpg'), |
|
'--project', 'out/', |
|
'--name', 'fixed_folder', |
|
'--exist-ok' |
|
] |
|
|
|
|
|
subprocess.run(command) |
|
|
|
|
|
output_image_path = 'out/fixed_folder/input_image_upscaled.jpg' |
|
|
|
|
|
if not os.path.exists(output_image_path): |
|
return "No output image found." |
|
|
|
|
|
output_image = cv2.imread(output_image_path) |
|
|
|
|
|
output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB) |
|
|
|
return output_image |
|
|
|
|
|
iface = gr.Interface( |
|
fn=detect_and_crop, |
|
inputs=gr.Image(type="pil"), |
|
outputs=gr.Image(type="numpy"), |
|
title="YOLOv7 Object Detection", |
|
description="Upload an image for object detection and cropping." |
|
) |
|
|
|
|
|
iface.launch() |