add inference size

app.py

@@ -74,6 +74,9 @@ def parse_args(known=False):
         type=str,
         help="cuda or cpu, hugging face only cpu",
     )
+    parser.add_argument(
+        "--inference_size", "-isz", default=640, type=int, help="model inference size"
+    )
 
     args = parser.parse_known_args()[0] if known else parser.parse_args()
     return args
@@ -115,7 +118,7 @@ def export_json(results, model, img_size):
 
 
 # YOLOv5图片检测函数
-def yolo_det(img, device, model_name, conf, iou, label_opt, model_cls):
+def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, model_cls):
 
     global model, model_name_tmp, device_tmp
 
@@ -133,7 +136,7 @@ def yolo_det(img, device, model_name, conf, iou, label_opt, model_cls):
     model.max_det = 1000  # 最大检测框数
     model.classes = model_cls  # 模型类别
 
-    results = model(img)  # 检测
+    results = model(img, size=inference_size)  # 检测
     results.render(labels=label_opt)  # 渲染
 
     det_img = Image.fromarray(results.imgs[0])  # 检测图片
@@ -178,6 +181,7 @@ def main(args):
     model_cfg = args.model_cfg
     cls_name = args.cls_name
     device = args.device
+    inference_size = args.inference_size
 
     # 模型加载
     model = model_loading(model_name, device)
@@ -193,6 +197,9 @@ def main(args):
     inputs_model = gr.inputs.Dropdown(
         choices=model_names, default=model_name, type="value", label="模型"
     )
+    inputs_size = gr.inputs.Radio(
+        choices=[320, 640], default=inference_size, label="推理尺寸"
+    )
     input_conf = gr.inputs.Slider(
         0, 1, step=slider_step, default=nms_conf, label="置信度阈值"
     )
@@ -209,6 +216,7 @@ def main(args):
         inputs_img,  # 输入图片
         device,  # 设备
         inputs_model,  # 模型
+        inputs_size,  # 推理尺寸
         input_conf,  # 置信度阈值
         inputs_iou,  # IoU阈值
         inputs_label,  # 标签显示
@@ -229,6 +237,7 @@ def main(args):
             "./img_example/bus.jpg",
             "cpu",
             "yolov5s",
+            640,
             0.6,
             0.5,
             True,
@@ -238,6 +247,7 @@ def main(args):
             "./img_example/Millenial-at-work.jpg",
             "cpu",
             "yolov5l",
+            320,
             0.5,
             0.45,
             True,
@@ -247,6 +257,7 @@ def main(args):
             "./img_example/zidane.jpg",
             "cpu",
             "yolov5m",
+            640,
             0.25,
             0.5,
             False,