v04 update

.gitignore

@@ -13,8 +13,12 @@
 
 # 日志格式
 *.log
-*.data
+*.datas
 *.txt
+
+# 生成文件
+*.pdf
+*.xlsx
 *.csv
 
 # 参数文件
@@ -32,6 +36,7 @@
 *.ttf
 *.otf
 
+# 模型文件
 *.pt
 *.db
 
@@ -41,5 +46,6 @@
 !cls_name/*
 !model_config/*
 !img_example/*
+!packages.txt
 
 app copy.py

app.py

@@ -133,16 +133,20 @@ def yaml_csv(file_path, file_tag):
 
 
 # model loading
-def model_loading(model_name, device, opt=["refresh_yolov5"]):
-
-    # load model
-    model = torch.hub.load(
-        model_path, 
-        model_name, 
-        force_reload=[True if "refresh_yolov5" in opt else False][0], 
-        device=device, 
-        _verbose=False
-    )
+def model_loading(model_name, device, opt=[]):
+
+    # 加载本地模型
+    try:
+        # load model
+        model = torch.hub.load(model_path,
+                               model_name,
+                               force_reload=[True if "refresh_yolov5" in opt else False][0],
+                               device=device,
+                               _verbose=False)
+    except Exception as e:
+        print(e)
+    else:
+        print(f"🚀 welcome to {GYD_VERSION}，{model_name} loaded successfully!")
 
     return model
 
@@ -174,13 +178,13 @@ def pil_draw(img, countdown_msg, textFont, xyxy, font_size, opt, obj_cls_index,
 
     if "label" in opt:
         text_w, text_h = textFont.getsize(countdown_msg)  # Label size
-        
+
         img_pil.rectangle(
             (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),
             fill=color_list[obj_cls_index],
             outline=color_list[obj_cls_index],
         )  # label background
-        
+
         img_pil.multiline_text(
             (xyxy[0], xyxy[1]),
             countdown_msg,
@@ -199,7 +203,7 @@ def color_set(cls_num):
         color = tuple(np.random.choice(range(256), size=3))
         # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
         color_list.append(color)
-    
+
     return color_list
 
 
@@ -218,9 +222,15 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
     if model_name_tmp != model_name:
         # Model judgment to avoid repeated loading
         model_name_tmp = model_name
+        print(f"Loading model {model_name_tmp}......")
         model = model_loading(model_name_tmp, device, opt)
     elif device_tmp != device:
+        # Device judgment to avoid repeated loading
         device_tmp = device
+        print(f"Loading model {model_name_tmp}......")
+        model = model_loading(model_name_tmp, device, opt)
+    else:
+        print(f"Loading model {model_name_tmp}......")
         model = model_loading(model_name_tmp, device, opt)
 
     # -------------Model tuning -------------
@@ -229,12 +239,12 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
     model.max_det = int(max_num)  # Maximum number of detection frames
     model.classes = model_cls  # model classes
 
-    color_list = color_set(len(model_cls_name_cp)) # 设置颜色
+    color_list = color_set(len(model_cls_name_cp))  # 设置颜色
 
     img_size = img.size  # frame size
 
     results = model(img, size=infer_size)  # detection
-    
+
     # ----------------目标裁剪----------------
     crops = results.crop(save=False)
     img_crops = []
@@ -243,7 +253,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
 
     # Data Frame
     dataframe = results.pandas().xyxy[0].round(2)
-    
+
     det_csv = "./Det_Report.csv"
     det_excel = "./Det_Report.xlsx"
 
@@ -251,7 +261,7 @@ def yolo_det_img(img, device, model_name, infer_size, conf, iou, max_num, model_
         dataframe.to_csv(det_csv, index=False)
     else:
         det_csv = None
-    
+
     if "excel" in opt:
         dataframe.to_excel(det_excel, sheet_name='sheet1', index=False)
     else:
@@ -363,9 +373,15 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
     if model_name_tmp != model_name:
         # Model judgment to avoid repeated loading
         model_name_tmp = model_name
+        print(f"Loading model {model_name_tmp}......")
         model = model_loading(model_name_tmp, device, opt)
     elif device_tmp != device:
+        # Device judgment to avoid repeated loading
         device_tmp = device
+        print(f"Loading model {model_name_tmp}......")
+        model = model_loading(model_name_tmp, device, opt)
+    else:
+        print(f"Loading model {model_name_tmp}......")
         model = model_loading(model_name_tmp, device, opt)
 
     # -------------Model tuning -------------
@@ -374,7 +390,7 @@ def yolo_det_video(video, device, model_name, infer_size, conf, iou, max_num, mo
     model.max_det = int(max_num)  # Maximum number of detection frames
     model.classes = model_cls  # model classes
 
-    color_list = color_set(len(model_cls_name_cp)) # 设置颜色
+    color_list = color_set(len(model_cls_name_cp))  # 设置颜色
 
     # ----------------Load fonts----------------
     yaml_index = cls_name.index(".yaml")
@@ -551,7 +567,9 @@ def main(args):
     outputs_video = gr.Video(format='mp4', label="Detection video")
 
     # output parameters
-    outputs_img_list = [outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf, outputs_csv, outputs_excel]
+    outputs_img_list = [
+        outputs_img, outputs_crops, outputs_objSize, outputs_clsSize, outputs_df, outputs_json, outputs_pdf,
+        outputs_csv, outputs_excel]
     outputs_video_list = [outputs_video]
 
     # title

packages.txt

@@ -0,0 +1,3 @@
+ffmpeg
+x264
+libx264-dev