hashb's picture
Update app.py
1e1f9c8
import cv2
import gradio as gr
from PIL import Image
import numpy as np
with open('coco.names', 'r') as f:
classes = f.read().splitlines()
net = cv2.dnn.readNetFromDarknet('yolov4.cfg', 'yolov4.weights')
model = cv2.dnn_DetectionModel(net)
model.setInputParams(scale=1 / 255, size=(416, 416), swapRB=True)
def detect(image):
img = Image.fromarray(image)
img = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
classIds, scores, boxes = model.detect(img, confThreshold=0.6, nmsThreshold=0.4)
for (classId, score, box) in zip(classIds, scores, boxes):
cv2.rectangle(img, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]),
color=(0, 255, 0), thickness=2)
text = '%s: %.2f' % (classes[classId], score)
cv2.putText(img, text, (box[0], box[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1,
color=(0, 255, 0), thickness=2)
return img
image_in = gr.Image()
image_out = gr.components.Image()
Iface = gr.Interface(
fn=detect,
inputs=image_in,
outputs=image_out,
title="Object Detection with YOLOS",
).launch()