Zengyf-CVer commited on
Commit
9fca2c9
1 Parent(s): a75f393

v04 update

Browse files
Files changed (2) hide show
  1. .gitignore +6 -1
  2. app.py +17 -17
.gitignore CHANGED
@@ -13,8 +13,12 @@
13
 
14
  # 日志格式
15
  *.log
16
- *.data
17
  *.txt
 
 
 
 
18
  *.csv
19
 
20
  # 参数文件
@@ -32,6 +36,7 @@
32
  *.ttf
33
  *.otf
34
 
 
35
  *.pt
36
  *.db
37
 
 
13
 
14
  # 日志格式
15
  *.log
16
+ *.datas
17
  *.txt
18
+
19
+ # 生成文件
20
+ *.pdf
21
+ *.xlsx
22
  *.csv
23
 
24
  # 参数文件
 
36
  *.ttf
37
  *.otf
38
 
39
+ # 模型文件
40
  *.pt
41
  *.db
42
 
app.py CHANGED
@@ -138,13 +138,11 @@ def model_loading(model_name, device, opt=[]):
138
  # 加载本地模型
139
  try:
140
  # load model
141
- model = torch.hub.load(
142
- model_path,
143
- model_name,
144
- force_reload=[True if "refresh_yolov5" in opt else False][0],
145
- device=device,
146
- _verbose=False
147
- )
148
  except Exception as e:
149
  print(e)
150
  else:
@@ -180,13 +178,13 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt, obj_cls_index,
180
 
181
  if "label" in opt:
182
  text_w, text_h = textFont.getsize(countdown_msg) # Label size
183
-
184
  img_pil.rectangle(
185
  (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),
186
  fill=color_list[obj_cls_index],
187
  outline=color_list[obj_cls_index],
188
  ) # label background
189
-
190
  img_pil.multiline_text(
191
  (xyxy[0], xyxy[1]),
192
  countdown_msg,
@@ -205,7 +203,7 @@ def color_set(cls_num):
205
  color = tuple(np.random.choice(range(256), size=3))
206
  # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
207
  color_list.append(color)
208
-
209
  return color_list
210
 
211
 
@@ -241,12 +239,12 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
241
  model.max_det = int(max_num) # Maximum number of detection frames
242
  model.classes = model_cls # model classes
243
 
244
- color_list = color_set(len(model_cls_name_cp)) # 设置颜色
245
 
246
  img_size = img.size # frame size
247
 
248
  results = model(img, size=infer_size) # detection
249
-
250
  # ----------------目标裁剪----------------
251
  crops = results.crop(save=False)
252
  img_crops = []
@@ -255,7 +253,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
255
 
256
  # Data Frame
257
  dataframe = results.pandas().xyxy[0].round(2)
258
-
259
  det_csv = "./Det_Report.csv"
260
  det_excel = "./Det_Report.xlsx"
261
 
@@ -263,7 +261,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
263
  dataframe.to_csv(det_csv, index=False)
264
  else:
265
  det_csv = None
266
-
267
  if "excel" in opt:
268
  dataframe.to_excel(det_excel, sheet_name='sheet1', index=False)
269
  else:
@@ -385,14 +383,14 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
385
  else:
386
  print(f"Loading model {model_name_tmp}......")
387
  model = model_loading(model_name_tmp, device, opt)
388
-
389
  # -------------Model tuning -------------
390
  model.conf = conf # NMS confidence threshold
391
  model.iou = iou # NMS IOU threshold
392
  model.max_det = int(max_num) # Maximum number of detection frames
393
  model.classes = model_cls # model classes
394
 
395
- color_list = color_set(len(model_cls_name_cp)) # 设置颜色
396
 
397
  # ----------------Load fonts----------------
398
  yaml_index = cls_name.index(".yaml")
@@ -569,7 +567,9 @@ def main(args):
569
  outputs_video = gr.Video(format='mp4', label="Detection video")
570
 
571
  # output parameters
572
- outputs_img_list = [outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf, outputs_csv, outputs_excel]
 
 
573
  outputs_video_list = [outputs_video]
574
 
575
  # title
 
138
  # 加载本地模型
139
  try:
140
  # load model
141
+ model = torch.hub.load(model_path,
142
+ model_name,
143
+ force_reload=[True if "refresh_yolov5" in opt else False][0],
144
+ device=device,
145
+ _verbose=False)
 
 
146
  except Exception as e:
147
  print(e)
148
  else:
 
178
 
179
  if "label" in opt:
180
  text_w, text_h = textFont.getsize(countdown_msg) # Label size
181
+
182
  img_pil.rectangle(
183
  (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),
184
  fill=color_list[obj_cls_index],
185
  outline=color_list[obj_cls_index],
186
  ) # label background
187
+
188
  img_pil.multiline_text(
189
  (xyxy[0], xyxy[1]),
190
  countdown_msg,
 
203
  color = tuple(np.random.choice(range(256), size=3))
204
  # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
205
  color_list.append(color)
206
+
207
  return color_list
208
 
209
 
 
239
  model.max_det = int(max_num) # Maximum number of detection frames
240
  model.classes = model_cls # model classes
241
 
242
+ color_list = color_set(len(model_cls_name_cp)) # 设置颜色
243
 
244
  img_size = img.size # frame size
245
 
246
  results = model(img, size=infer_size) # detection
247
+
248
  # ----------------目标裁剪----------------
249
  crops = results.crop(save=False)
250
  img_crops = []
 
253
 
254
  # Data Frame
255
  dataframe = results.pandas().xyxy[0].round(2)
256
+
257
  det_csv = "./Det_Report.csv"
258
  det_excel = "./Det_Report.xlsx"
259
 
 
261
  dataframe.to_csv(det_csv, index=False)
262
  else:
263
  det_csv = None
264
+
265
  if "excel" in opt:
266
  dataframe.to_excel(det_excel, sheet_name='sheet1', index=False)
267
  else:
 
383
  else:
384
  print(f"Loading model {model_name_tmp}......")
385
  model = model_loading(model_name_tmp, device, opt)
386
+
387
  # -------------Model tuning -------------
388
  model.conf = conf # NMS confidence threshold
389
  model.iou = iou # NMS IOU threshold
390
  model.max_det = int(max_num) # Maximum number of detection frames
391
  model.classes = model_cls # model classes
392
 
393
+ color_list = color_set(len(model_cls_name_cp)) # 设置颜色
394
 
395
  # ----------------Load fonts----------------
396
  yaml_index = cls_name.index(".yaml")
 
567
  outputs_video = gr.Video(format='mp4', label="Detection video")
568
 
569
  # output parameters
570
+ outputs_img_list = [
571
+ outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf,
572
+ outputs_csv, outputs_excel]
573
  outputs_video_list = [outputs_video]
574
 
575
  # title