Zengyf-CVer commited on
Commit
c973154
1 Parent(s): 80cc235

v04 update

Browse files
Files changed (1) hide show
  1. app.py +16 -9
app.py CHANGED
@@ -133,10 +133,16 @@ def yaml_csv(file_path, file_tag):
133
 
134
 
135
  # model loading
136
- def model_loading(model_name, device):
137
 
138
  # load model
139
- model = torch.hub.load(model_path, model_name, force_reload=True, device=device, _verbose=False)
 
 
 
 
 
 
140
 
141
  return model
142
 
@@ -186,6 +192,7 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt, obj_cls_index,
186
  return img
187
 
188
 
 
189
  def color_set(cls_num):
190
  color_list = []
191
  for i in range(cls_num):
@@ -211,10 +218,10 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
211
  if model_name_tmp != model_name:
212
  # Model judgment to avoid repeated loading
213
  model_name_tmp = model_name
214
- model = model_loading(model_name_tmp, device)
215
  elif device_tmp != device:
216
  device_tmp = device
217
- model = model_loading(model_name_tmp, device)
218
 
219
  # -------------Model tuning -------------
220
  model.conf = conf # NMS confidence threshold
@@ -343,10 +350,10 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
343
  if model_name_tmp != model_name:
344
  # Model judgment to avoid repeated loading
345
  model_name_tmp = model_name
346
- model = model_loading(model_name_tmp, device)
347
  elif device_tmp != device:
348
  device_tmp = device
349
- model = model_loading(model_name_tmp, device)
350
 
351
  # -------------Model tuning -------------
352
  model.conf = conf # NMS confidence threshold
@@ -472,7 +479,7 @@ def main(args):
472
  inputs_iou01 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU threshold")
473
  inputs_maxnum01 = gr.Number(value=max_detnum, label="Maximum number of detections")
474
  inputs_clsName01 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="category")
475
- inputs_opt01 = gr.CheckboxGroup(choices=["label", "pdf", "json"],
476
  value=["label", "pdf"],
477
  type="value",
478
  label="operate")
@@ -486,7 +493,7 @@ def main(args):
486
  inputs_iou02 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU threshold")
487
  inputs_maxnum02 = gr.Number(value=max_detnum, label="Maximum number of detections")
488
  inputs_clsName02 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="category")
489
- inputs_opt02 = gr.CheckboxGroup(choices=["label"], value=["label"], type="value", label="operate")
490
 
491
  # Input parameters
492
  inputs_img_list = [
@@ -536,7 +543,7 @@ def main(args):
536
  title = "Gradio YOLOv5 Det v0.4"
537
 
538
  # describe
539
- description = "<div align='center'>Customizable target detection model, easy to install, easy to use</div>"
540
  # article="https://gitee.com/CV_Lab/gradio_yolov5_det"
541
 
542
  # example image
 
133
 
134
 
135
  # model loading
136
+ def model_loading(model_name, device, opt=["refresh_yolov5"]):
137
 
138
  # load model
139
+ model = torch.hub.load(
140
+ model_path,
141
+ model_name,
142
+ force_reload=[True if "refresh_yolov5" in opt else False][0],
143
+ device=device,
144
+ _verbose=False
145
+ )
146
 
147
  return model
148
 
 
192
  return img
193
 
194
 
195
+ # Label and bounding box color settings
196
  def color_set(cls_num):
197
  color_list = []
198
  for i in range(cls_num):
 
218
  if model_name_tmp != model_name:
219
  # Model judgment to avoid repeated loading
220
  model_name_tmp = model_name
221
+ model = model_loading(model_name_tmp, device, opt)
222
  elif device_tmp != device:
223
  device_tmp = device
224
+ model = model_loading(model_name_tmp, device, opt)
225
 
226
  # -------------Model tuning -------------
227
  model.conf = conf # NMS confidence threshold
 
350
  if model_name_tmp != model_name:
351
  # Model judgment to avoid repeated loading
352
  model_name_tmp = model_name
353
+ model = model_loading(model_name_tmp, device, opt)
354
  elif device_tmp != device:
355
  device_tmp = device
356
+ model = model_loading(model_name_tmp, device, opt)
357
 
358
  # -------------Model tuning -------------
359
  model.conf = conf # NMS confidence threshold
 
479
  inputs_iou01 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU threshold")
480
  inputs_maxnum01 = gr.Number(value=max_detnum, label="Maximum number of detections")
481
  inputs_clsName01 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="category")
482
+ inputs_opt01 = gr.CheckboxGroup(choices=["label", "pdf", "json", "refresh_yolov5"],
483
  value=["label", "pdf"],
484
  type="value",
485
  label="operate")
 
493
  inputs_iou02 = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU threshold")
494
  inputs_maxnum02 = gr.Number(value=max_detnum, label="Maximum number of detections")
495
  inputs_clsName02 = gr.CheckboxGroup(choices=model_cls_name, value=model_cls_name, type="index", label="category")
496
+ inputs_opt02 = gr.CheckboxGroup(choices=["label", "refresh_yolov5"], value=["label"], type="value", label="operate")
497
 
498
  # Input parameters
499
  inputs_img_list = [
 
543
  title = "Gradio YOLOv5 Det v0.4"
544
 
545
  # describe
546
+ description = "Author: 曾逸夫(Zeng Yifu), Project Address: https://gitee.com/CV_Lab/gradio_yolov5_det, Github: https://github.com/Zengyf-CVer, thanks to [Gradio](https://github.com/gradio-app/gradio) & [YOLOv5](https://github.com/ultralytics/yolov5)"
547
  # article="https://gitee.com/CV_Lab/gradio_yolov5_det"
548
 
549
  # example image