|
import gradio as gr |
|
import matplotlib.pyplot as plt |
|
from PIL import Image |
|
from ultralytics import YOLO |
|
import cv2 |
|
import numpy as np |
|
|
|
def image_preprocess(image): |
|
|
|
img_height, img_width = image.shape[0:2] |
|
image_converted = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|
ih, iw = [input_size, input_size] |
|
h, w, _ = image.shape |
|
|
|
scale = min(iw/w, ih/h) |
|
nw, nh = int(scale * w), int(scale * h) |
|
image_resized = cv2.resize(image_converted, (nw, nh)) |
|
|
|
image_padded = np.full(shape=[ih, iw, 3], fill_value=128.0) |
|
dw, dh = (iw - nw) // 2, (ih-nh) // 2 |
|
image_padded[dh:nh+dh, dw:nw+dw, :] = image_resized |
|
image_padded = image_padded / 255. |
|
|
|
image_padded = image_padded[np.newaxis, ...].astype(np.float32) |
|
image_padded = np.moveaxis(image_padded, -1, 1) |
|
|
|
|
|
return image_padded, img_width, img_height, image |
|
|
|
model = YOLO('best (1).pt') |
|
|
|
def response(image): |
|
results = model(image) |
|
for i, r in enumerate(results): |
|
|
|
im_bgr = r.plot() |
|
im_rgb = im_bgr[..., ::-1] |
|
|
|
|
|
return im_rgb |
|
|
|
iface = gr.Interface(fn=response, inputs="image", outputs="image") |
|
iface.launch() |