Zengyf-CVer commited on
Commit
53f1c15
1 Parent(s): e351b90

v04 update

Browse files
Files changed (3) hide show
  1. .gitignore +7 -1
  2. app.py +37 -19
  3. packages.txt +3 -0
.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
 
@@ -41,5 +46,6 @@
41
  !cls_name/*
42
  !model_config/*
43
  !img_example/*
 
44
 
45
  app copy.py
 
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
 
 
46
  !cls_name/*
47
  !model_config/*
48
  !img_example/*
49
+ !packages.txt
50
 
51
  app copy.py
app.py CHANGED
@@ -133,16 +133,20 @@ def yaml_csv(file_path, file_tag):
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
 
@@ -174,13 +178,13 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt, obj_cls_index,
174
 
175
  if "label" in opt:
176
  text_w, text_h = textFont.getsize(countdown_msg) # Label size
177
-
178
  img_pil.rectangle(
179
  (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),
180
  fill=color_list[obj_cls_index],
181
  outline=color_list[obj_cls_index],
182
  ) # label background
183
-
184
  img_pil.multiline_text(
185
  (xyxy[0], xyxy[1]),
186
  countdown_msg,
@@ -199,7 +203,7 @@ def color_set(cls_num):
199
  color = tuple(np.random.choice(range(256), size=3))
200
  # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
201
  color_list.append(color)
202
-
203
  return color_list
204
 
205
 
@@ -218,9 +222,15 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
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 -------------
@@ -229,12 +239,12 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
229
  model.max_det = int(max_num) # Maximum number of detection frames
230
  model.classes = model_cls # model classes
231
 
232
- color_list = color_set(len(model_cls_name_cp)) # 设置颜色
233
 
234
  img_size = img.size # frame size
235
 
236
  results = model(img, size=infer_size) # detection
237
-
238
  # ----------------目标裁剪----------------
239
  crops = results.crop(save=False)
240
  img_crops = []
@@ -243,7 +253,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
243
 
244
  # Data Frame
245
  dataframe = results.pandas().xyxy[0].round(2)
246
-
247
  det_csv = "./Det_Report.csv"
248
  det_excel = "./Det_Report.xlsx"
249
 
@@ -251,7 +261,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
251
  dataframe.to_csv(det_csv, index=False)
252
  else:
253
  det_csv = None
254
-
255
  if "excel" in opt:
256
  dataframe.to_excel(det_excel, sheet_name='sheet1', index=False)
257
  else:
@@ -363,9 +373,15 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
363
  if model_name_tmp != model_name:
364
  # Model judgment to avoid repeated loading
365
  model_name_tmp = model_name
 
366
  model = model_loading(model_name_tmp, device, opt)
367
  elif device_tmp != device:
 
368
  device_tmp = device
 
 
 
 
369
  model = model_loading(model_name_tmp, device, opt)
370
 
371
  # -------------Model tuning -------------
@@ -374,7 +390,7 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
374
  model.max_det = int(max_num) # Maximum number of detection frames
375
  model.classes = model_cls # model classes
376
 
377
- color_list = color_set(len(model_cls_name_cp)) # 设置颜色
378
 
379
  # ----------------Load fonts----------------
380
  yaml_index = cls_name.index(".yaml")
@@ -551,7 +567,9 @@ def main(args):
551
  outputs_video = gr.Video(format='mp4', label="Detection video")
552
 
553
  # output parameters
554
- outputs_img_list = [outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf, outputs_csv, outputs_excel]
 
 
555
  outputs_video_list = [outputs_video]
556
 
557
  # title
 
133
 
134
 
135
  # model loading
136
+ def model_loading(model_name, device, opt=[]):
137
+
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:
149
+ print(f"🚀 welcome to {GYD_VERSION},{model_name} loaded successfully!")
150
 
151
  return model
152
 
 
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
 
 
222
  if model_name_tmp != model_name:
223
  # Model judgment to avoid repeated loading
224
  model_name_tmp = model_name
225
+ print(f"Loading model {model_name_tmp}......")
226
  model = model_loading(model_name_tmp, device, opt)
227
  elif device_tmp != device:
228
+ # Device judgment to avoid repeated loading
229
  device_tmp = device
230
+ print(f"Loading model {model_name_tmp}......")
231
+ model = model_loading(model_name_tmp, device, opt)
232
+ else:
233
+ print(f"Loading model {model_name_tmp}......")
234
  model = model_loading(model_name_tmp, device, opt)
235
 
236
  # -------------Model tuning -------------
 
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:
 
373
  if model_name_tmp != model_name:
374
  # Model judgment to avoid repeated loading
375
  model_name_tmp = model_name
376
+ print(f"Loading model {model_name_tmp}......")
377
  model = model_loading(model_name_tmp, device, opt)
378
  elif device_tmp != device:
379
+ # Device judgment to avoid repeated loading
380
  device_tmp = device
381
+ print(f"Loading model {model_name_tmp}......")
382
+ model = model_loading(model_name_tmp, device, opt)
383
+ else:
384
+ print(f"Loading model {model_name_tmp}......")
385
  model = model_loading(model_name_tmp, device, opt)
386
 
387
  # -------------Model tuning -------------
 
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
packages.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ ffmpeg
2
+ x264
3
+ libx264-dev