Update app.py

app.py

@@ -1,157 +1,11 @@
-import onnxruntime
-import numpy as np
-import cv2
-import os
 import gradio as gr
-from typing import List, Tuple
-
-img_size = 640
-
-class Object:
-    def __init__(self, rect: Tuple[int, int, int, int], label: int, prob: float):
-        self.rect = rect
-        self.label = label
-        self.prob = prob
-
-def sigmoid(x):
-    return 1/(1+np.exp(-x))
-
-def intersection_area(a: Object, b: Object) -> float:
-    x = max(a.rect[0], b.rect[0])
-    y = max(a.rect[1], b.rect[1])
-    w = min(a.rect[0] + a.rect[2], b.rect[0] + b.rect[2]) - x
-    h = min(a.rect[1] + a.rect[3], b.rect[1] + b.rect[3]) - y
-    if w < 0 or h < 0:
-        return 0
-    return w * h
-
-def nms_sorted_bboxes(faceobjects: List[Object], nms_threshold: float) -> List[int]:
-    picked = []
-    n = len(faceobjects)
-    areas = [faceobjects[i].rect[2]*faceobjects[i].rect[3] for i in range(n)]
-    for i in range(n):
-        a = faceobjects[i]
-        keep = 1
-        for j in range(len(picked)):
-            b = faceobjects[picked[j]]
-            inter_area = intersection_area(a, b)
-            union_area = areas[i] + areas[picked[j]] - inter_area
-            if inter_area / union_area > nms_threshold:
-                keep = 0
-        if keep:
-            picked.append(i)
-    return picked
-
-def generate_proposals(stride, feat, prob_threshold, letterbox_cols, letterbox_rows):
-    anchors = [10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326]
-    anchor_num = 3
-    feat_w = letterbox_cols // stride
-    feat_h = letterbox_rows // stride
-    cls_num = 45
-    anchor_group = 0
-    if stride == 8:
-        anchor_group = 1
-    if stride == 16:
-        anchor_group = 2
-    if stride == 32:
-        anchor_group = 3
-    objects = []
-    for h in range(feat_h):
-        for w in range(feat_w):
-            for a in range(anchor_num):
-                class_index = 0
-                class_score = float('-inf')
-                for s in range(cls_num):
-                    score = feat[a * feat_w * feat_h * (cls_num + 5) + h * feat_w * (cls_num + 5) + w * (cls_num + 5) + s + 5]
-                    # if score > class_score:
-                    if np.max(score) > class_score:
-                        class_index = s
-                        class_score = np.max(score)
-                box_score = feat[a * feat_w * feat_h * (cls_num + 5) + (h * feat_w) * (cls_num + 5) + w * (cls_num + 5) + 4]
-                final_score = sigmoid(box_score) * sigmoid(class_score)
-                if final_score >= prob_threshold:
-                    loc_idx = a * feat_h * feat_w * (cls_num + 5) + h * feat_w * (cls_num + 5) + w * (cls_num + 5)
-                    dx = sigmoid(feat[loc_idx + 0])
-                    dy = sigmoid(feat[loc_idx + 1])
-                    dw = sigmoid(feat[loc_idx + 2])
-                    dh = sigmoid(feat[loc_idx + 3])
-                    pred_cx = (dx * 2.0 - 0.5 + w) * stride
-                    pred_cy = (dy * 2.0 - 0.5 + h) * stride
-                    anchor_w = anchors[(anchor_group - 1) * 6 + a * 2 + 0]
-                    anchor_h = anchors[(anchor_group - 1) * 6 + a * 2 + 1]
-                    pred_w = dw * dw * 4.0 * anchor_w
-                    pred_h = dh * dh * 4.0 * anchor_h
-                    x0 = pred_cx - pred_w * 0.5
-                    y0 = pred_cy - pred_h * 0.5
-                    x1 = pred_cx + pred_w * 0.5
-                    y1 = pred_cy + pred_h * 0.5
-                    obj = Object((x0, y0, x1 - x0, y1 - y0), class_index, final_score)
-                    objects.append(obj)
-    return objects
-
-def draw_objects(bgr, objects):
-    class_names = ["pl80", "p6", "p5", "pm55", "pl60", "ip", "p11", "i2r", "p23", "pg",
-                   "il80", "ph4", "i4", "pl70", "pne", "ph4.5", "p12", "p3", "pl5", "w13",
-                   "i4l", "pl30", "p10", "pn", "w55", "p26", "p13", "pr40", "pl20", "pm30",
-                   "pl40", "i2", "pl120", "w32", "ph5", "il60", "w57", "pl100",
-                   "w59", "il100", "p19", "pm20", "i5", "p27", "pl50"]
-    image = bgr.copy()
-    for i, obj in enumerate(objects):
-        print(f"{obj.label}: {obj.prob*100}%, {[obj.rect[0], obj.rect[1], obj.rect[0]+obj.rect[2], obj.rect[1]+obj.rect[3]]}, {class_names[obj.label]}")
-        cv2.rectangle(image, (int(obj.rect[0]), int(obj.rect[1])), (int(obj.rect[0]+obj.rect[2]), int(obj.rect[1]+obj.rect[3])), (255, 0, 0), 2)
-        label = f"{class_names[obj.label]} {obj.prob*100:.1f}%"
-        label_size, base_line = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(image, (int(obj.rect[0]), int(obj.rect[1]-label_size[1])), (int(obj.rect[0])+label_size[0], int(obj.rect[1]+base_line)), (255, 255, 255), cv2.FILLED)
-        cv2.putText(image, label, (int(obj.rect[0]), int(obj.rect[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    return image
-
-def preprocess(image_file: str, img_size: int = 640):
-    img = cv2.imread(image_file)
-    h, w, _ = img.shape
-    r = img_size / max(h, w)
-    nh, nw = int(r * h), int(r * w)
-    img = cv2.resize(img, (nw, nh))
-    img = cv2.copyMakeBorder(img, (img_size - nh) // 2, (img_size - nh) // 2,
-                             (img_size - nw) // 2, (img_size - nw) // 2,
-                             cv2.BORDER_CONSTANT, value=(114, 114, 114))
-    img = img.astype(np.float32) / 255.0
-    img = np.transpose(img, [2, 0, 1])
-    img = np.expand_dims(img, axis=0)
-    return img
-
-def infer_main(image_file: str):
-    model_file = "./jtbz_best.onnx"
-
-    # load model
-    sess = onnxruntime.InferenceSession(model_file)
-    input_name = sess.get_inputs()[0].name
-
-    # preprocess
-    img = preprocess(image_file)
-
-    # inference
-    outputs = sess.run(None, {input_name: img})
-
-    # postprocess
-    prob_threshold = 0.27
-    nms_threshold = 0.45
-    proposals = []
-    for stride, feat in zip([8, 16, 32], outputs[:3]):
-        proposals.extend(generate_proposals(stride, feat[0], prob_threshold, img_size, img_size))
-    proposals.sort(key=lambda x: x.prob, reverse=True)
-    picked = nms_sorted_bboxes(proposals, nms_threshold)
-
-    # draw result
-    img = cv2.imread(image_file)
-    objects = [proposals[i] for i in picked]
-    result = draw_objects(img, objects)
+import subprocess
+import os
 
-    return result
+def run_yolov5s(image):
+    image.save("input.jpg")
+    result = subprocess.run(["./yolov5s-tt100k", "-m", "./yolov5s_tt100k_opt_fp32.tmfile", "-i", "input.jpg"], capture_output=True, text=True)
+    return "yolov5s-tt100k.out.jpg", result.stdout
 
-# gradio interface
-iface = gr.Interface(fn=infer_main, 
-                     inputs=gr.components.Image(type="filepath"), 
-                     outputs=gr.components.Image(),
-                     title="YOLOv5 Object Detection and classification",
-                     description="Upload an image to see the detected objects and their classification.")
-iface.launch()
+iface = gr.Interface(fn=run_yolov5s, inputs=gr.Image(), outputs=[gr.Image(type="filepath"), gr.Textbox()])
+iface.launch()