|
import os |
|
import gradio as gr |
|
import cv2 |
|
|
|
from ultralytics import YOLO |
|
|
|
""" |
|
input_dir = '/home/student/PycharmProjects/pythonProjectYoloObjectPredict/test_images' |
|
output_dir = '/home/student/PycharmProjects/pythonProjectYoloObjectPredict/' |
|
|
|
image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg')] |
|
|
|
for image_file in image_files: |
|
image_path = os.path.join(input_dir, image_file) |
|
output_path = os.path.join(output_dir, f'{os.path.splitext(image_file)[0]}_result.jpg') |
|
|
|
results = model.predict(image_path, conf=0.25, save=True, save_crop=True) |
|
print(f'Predictions for {image_file}: {results}') |
|
""" |
|
|
|
|
|
image_directory = '/content/drive/MyDrive/Work/yolo/images/val' |
|
|
|
jpg_files = [file for file in os.listdir(image_directory) if file.lower().endswith('.jpg')] |
|
|
|
|
|
path = [os.path.join(image_directory, filename) for filename in jpg_files] |
|
|
|
model = YOLO('/content/drive/MyDrive/Work/yolo/best.pt') |
|
|
|
inputs_image = [ |
|
gr.components.Image(type="filepath", label="Input Image"), |
|
] |
|
outputs_image = [ |
|
gr.components.Image(type="numpy", label="Output Image"), |
|
] |
|
|
|
def show_preds_image(image_path): |
|
image = cv2.imread(image_path) |
|
outputs = model.predict(source=image_path) |
|
results = outputs[0].cpu().numpy() |
|
for i, det in enumerate(results.boxes.xyxy): |
|
cv2.rectangle( |
|
image, |
|
(int(det[0]), int(det[1])), |
|
(int(det[2]), int(det[3])), |
|
color=(0, 0, 255), |
|
thickness=2, |
|
lineType=cv2.LINE_AA |
|
) |
|
return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|
|
|
interface_image = gr.Interface( |
|
fn=show_preds_image, |
|
inputs=inputs_image, |
|
outputs=outputs_image, |
|
title="Floor Plan Detector", |
|
examples=path, |
|
cache_examples=False, |
|
) |
|
|
|
gr.TabbedInterface( |
|
[interface_image], |
|
tab_names=['Image Inference'], |
|
).queue().launch(debug=True) |