add model auto downloads

app.py

@@ -16,8 +16,9 @@ from ultralytics import YOLO
 ROOT_PATH = sys.path[0]  # 项目根目录
 
 fonts_list = ["SimSun.ttf", "TimesNewRoman.ttf", "malgun.ttf"]  # 字体列表
+models_list = ["cnn_se.pt", "detr_based.pt", "vit_based.pt", "yolov5_based.pt", "yolov8_based.pt"] # 模型列表
 fonts_directory_path = Path(ROOT_PATH, "fonts")
-models_directory_path = Path(ROOT_PATH)  # 模型存放在项目的根目录
+models_directory_path = Path(ROOT_PATH, "models")  # 模型存放在项目的根目录
 
 data_url_dict = {
     "SimSun.ttf": "https://raw.githubusercontent.com/Tsumugii24/Typora-images/main/files/SimSun.ttf",
@@ -37,9 +38,9 @@ model_url_dict = {
 def is_fonts(fonts_dir):
     if fonts_dir.is_dir():
         # 如果本地字体库存在
-        local_list = os.listdir(fonts_dir)  # 本地字体库
+        local_font_list = os.listdir(fonts_dir)  # 本地字体库
 
-        font_diff = list(set(fonts_list).difference(set(local_list)))
+        font_diff = list(set(fonts_list).difference(set(local_font_list)))
 
         if font_diff != []:
             # 缺失字体
@@ -55,19 +56,19 @@ def is_fonts(fonts_dir):
 def is_models(models_dir):
     if models_dir.is_dir():
         # 如果本地模型库存在
-        local_list = os.listdir(models_dir)  # 本地模型库
+        local_model_list = os.listdir(models_dir)  # 本地模型库
 
-        model_diff = list(set(model_url_dict.keys()).difference(set(local_list)))
+        model_diff = list(set(models_list()).difference(set(local_model_list)))
 
         if model_diff != []:
             # 缺失模型
             download_models(model_diff)  # 下载缺失的模型
         else:
-            print(f"{model_url_dict.keys()}[bold green]Required models already downloaded！[/bold green]")
+            print(f"{models_list}[bold green]Required models already downloaded！[/bold green]")
     else:
         # 本地模型库不存在，创建模型库
         print("[bold red]Local models library does not exist, creating now...[/bold red]")
-        download_models(model_url_dict.keys())  # 创建模型库
+        download_models(models_list)  # 创建模型库
 
 # 下载字体
 def download_fonts(font_diff):
@@ -78,9 +79,9 @@ def download_fonts(font_diff):
             font_name = v.split("/")[-1]  # 字体名称
             fonts_directory_path.mkdir(parents=True, exist_ok=True)  # 创建目录
 
-            file_path = f"{ROOT_PATH}/fonts/{font_name}"  # 字体路径
+            font_file_path = f"{ROOT_PATH}/fonts/{font_name}"  # 字体路径
             # 下载字体文件
-            wget.download(v, file_path)
+            wget.download(v, font_file_path)
 
 # 下载模型
 def download_models(model_diff):
@@ -90,9 +91,9 @@ def download_models(model_diff):
         v = model_url_dict[k]
         model_name = v.split("/")[-1]  # 模型名称
 
-        file_path = f"{ROOT_PATH}/{model_name}"  # 模型路径
+        model_file_path = f"{ROOT_PATH}/models/{model_name}"  # 模型路径
         # 下载模型文件
-        wget.download(v, file_path)
+        wget.download(v, model_file_path)
 
 
 is_fonts(fonts_directory_path)
@@ -336,7 +337,7 @@ def seg_output(img_path, seg_mask_list, color_list, cls_list):
 
 
 # 目标检测和图像分割模型加载
-def model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"):
+def model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8_based.pt"):
     model = YOLO(yolo_model)
 
     results = model(source=img_path, device=device_opt, imgsz=infer_size, conf=conf, iou=iou, max_det=max_det)
@@ -358,7 +359,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
     cls_index_det_stat = []  # 1
 
     # 模型加载
-    predict_results = model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model=f"{model_name}.pt")
+    predict_results = model_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model=f"models/{model_name}.pt")
     # 检测参数
     xyxy_list = predict_results.boxes.xyxy.cpu().numpy().tolist()
     conf_list = predict_results.boxes.conf.cpu().numpy().tolist()