Dricz's picture
Update app.py
b0e1d53 verified
raw
history blame
1.21 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):
res = image_preprocess(image)
im_rgb = model(res)
return im_rgb
iface = gr.Interface(fn=response, inputs="image", outputs="image")
iface.launch()