Dricz's picture
Update app.py
c134122 verified
raw
history blame
No virus
1.41 kB
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] # [input_size, input_size] = [640, 640]
h, w, _ = image.shape # [1944, 2592]
scale = min(iw/w, ih/h) # min(0.2469, 0.3292) = 0.2469
nw, nh = int(scale * w), int(scale * h) # [640, 480]
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 # [0, 80]
image_padded[dh:nh+dh, dw:nw+dw, :] = image_resized # image_padded[80:256, 32:224]
image_padded = image_padded / 255.
# image_resized = image_resized / 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):
print(image)
results = model(image)
for i, r in enumerate(results):
# Plot results image
im_bgr = r.plot() # BGR-order numpy array
im_rgb = im_bgr[..., ::-1] # Convert BGR to RGB
# im_rgb = Image.fromarray(im_rgb)
return im_rgb
iface = gr.Interface(fn=response, inputs="image", outputs="image")
iface.launch()