add model auto downloads

app.py

@@ -17,6 +17,7 @@ ROOT_PATH = sys.path[0]  # 项目根目录
 
 fonts_list = ["SimSun.ttf", "TimesNewRoman.ttf", "malgun.ttf"]  # 字体列表
 fonts_directory_path = Path(ROOT_PATH, "fonts")
+models_directory_path = Path(ROOT_PATH)  # 模型存放在项目的根目录
 
 data_url_dict = {
     "SimSun.ttf": "https://raw.githubusercontent.com/Tsumugii24/Typora-images/main/files/SimSun.ttf",
@@ -24,6 +25,13 @@ data_url_dict = {
     "malgun.ttf": "https://raw.githubusercontent.com/Tsumugii24/Typora-images/main/files/malgun.ttf",
 }
 
+model_url_dict = {
+    "cnn_se.pt": "https://huggingface.co/Tsumugii/lesion-cells-det/raw/main/cnn_se.pt",
+    "detr_based.pt": "https://huggingface.co/Tsumugii/lesion-cells-det/raw/main/detr_based.pt",
+    "vit_based.pt": "https://huggingface.co/Tsumugii/lesion-cells-det/raw/main/vit_based.pt",
+    "yolov5_based.pt": "https://huggingface.co/Tsumugii/lesion-cells-det/raw/main/yolov5_based.pt",
+    "yolov8_based.pt": "https://huggingface.co/Tsumugii/lesion-cells-det/raw/main/yolov8_based.pt",
+}
 
 # 判断字体文件是否存在
 def is_fonts(fonts_dir):
@@ -42,6 +50,24 @@ def is_fonts(fonts_dir):
         # 本地字体库不存在，创建字体库
         print("[bold red]Local fonts library does not exist, creating now...[/bold red]")
         download_fonts(fonts_list)  # 创建字体库
+        
+# 判断模型文件是否存在
+def is_models(models_dir):
+    if models_dir.is_dir():
+        # 如果本地模型库存在
+        local_list = os.listdir(models_dir)  # 本地模型库
+
+        model_diff = list(set(model_url_dict.keys()).difference(set(local_list)))
+
+        if model_diff != []:
+            # 缺失模型
+            download_models(model_diff)  # 下载缺失的模型
+        else:
+            print(f"{model_url_dict.keys()}[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())  # 创建模型库
 
 # 下载字体
 def download_fonts(font_diff):
@@ -56,7 +82,21 @@ def download_fonts(font_diff):
             # 下载字体文件
             wget.download(v, file_path)
 
+# 下载模型
+def download_models(model_diff):
+    global model_name
+
+    for k in model_diff:
+        v = model_url_dict[k]
+        model_name = v.split("/")[-1]  # 模型名称
+
+        file_path = f"{ROOT_PATH}/{model_name}"  # 模型路径
+        # 下载模型文件
+        wget.download(v, file_path)
+
+
 is_fonts(fonts_directory_path)
+is_models(models_directory_path)
 
 # --------------------- 字体库 ---------------------
 SimSun_path = f"{ROOT_PATH}/fonts/SimSun.ttf"  # 宋体文件路径