import cv2 as cv
import sys
import numpy as np
import gradio as gr
sys.path.append('../')
from include.SoyNet import *
from utils.ClassName import COCO_80
class_names = COCO_80()
# gradio interface에서 사용할 fucntion 정의
def predict(img):
resized_img = cv.resize(img, (input_width, input_height))
# 추론결과 output 데이터 형식정의
data_type = np.dtype([("x1", c_float), ("y1", c_float), ("x2", c_float), ("y2", c_float),
("obj_id", c_int), ("prob", c_float)])
output = np.zeros(batch_size * nms_count, dtype=data_type)
# SoyNet을 이용한 객체감지 추론
feedData(handle, resized_img)
inference(handle)
getOutput(handle, output)
# 감지결과 화면 표시
threshold = 0.6
for n_idx in range(nms_count):
x1, y1, x2, y2, obj_id, prob = output[n_idx]
if prob > threshold:
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cv.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv.putText(img, class_names[obj_id], (x1, y1 - 3), 1, 1.5, (255, 0, 0), 1, cv.LINE_AA)
return img
if __name__ == "__main__":
# SoyNet 핸들 생성을 위한 초기값 설정
engine_serialize = 0
batch_size = 1
region_count = 1000
nms_count = 100
class_count = len(class_names)
input_height, input_width = 720, 1280
model_size = 640
model_name = "yolov5s"
cfg_file = "../models/Yolov5-6.1-s/configs/{}.cfg".format(model_name)
weight_file = "../models/Yolov5-6.1-s/weights/{}.weights".format(model_name)
engine_file = "../models/Yolov5-6.1-s/engines/{}.bin".format(model_name)
log_file = "../models/Yolov5-6.1-s/logs/{}.log".format(model_name)
extend_param = \
"MODEL_NAME={} BATCH_SIZE={} ENGINE_SERIALIZE={} CLASS_COUNT={} NMS_COUNT={} REGION_COUNT={} " \
"INPUT_SIZE={},{} MODEL_SIZE={},{} " \
"WEIGHT_FILE={} ENGINE_FILE={} LOG_FILE={}".format(
model_name, batch_size, engine_serialize, class_count, nms_count, region_count,
input_height, input_width, model_size, model_size,
weight_file, engine_file, log_file
)
# SoyNet 핸들 생성
handle = initSoyNet(cfg_file, extend_param)
# Gradio Interface 생성
demo = gr.Interface(
fn=predict,
inputs="image",
outputs="image",
title='SoyNet gradio Demo',
description='This is gradio demo using SoyNet inference engine.',
examples=['examples/NY_01.png','examples/NY_02.png'],
)
try:
demo.launch()
except KeyboardInterrupt:
# SoyNet 핸들 해제
freeSoyNet(handle) | 
