|
import os |
|
import gradio as gr |
|
import cv2 |
|
|
|
from ultralytics import YOLO |
|
|
|
image_directory = '/home/user/app/some_flat_plans' |
|
|
|
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('/home/user/app/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) |