Dricz's picture
Update app.py
5d64bc0 verified
raw
history blame
1.25 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.fromarray(image)
im_rgb = model(res)
# im_rgb = Image.fromarray(im_rgb)
return np.array(im_rgb)
iface = gr.Interface(fn=response, inputs="image", outputs="image")
iface.launch()