Update app.py

app.py

@@ -11,7 +11,7 @@ import torch.autograd as autograd
 import numpy as np
 import os
 from datetime import datetime
-
+from huggingface_hub import hf_hub_download
 def creatMask(Image, threshold=5):
     ##This program try to creat the mask for the filed-of-view
     ##Input original image (RGB or green channel), threshold (user set parameter, default 10)
@@ -279,10 +279,9 @@ def GetResult_out_big(Net, k, use_cuda=False, dataset='', is_kill_border=False,
     return ArteryPred, VeinPred, VesselPred, Mask, ArteryPred, VeinPred, VesselPred,
 
 
-def out_test(cfg, output_dir='', evaluate_metrics=False, img_name='out_test'):
+def out_test(cfg,model_path='', output_dir='', evaluate_metrics=False, img_name='out_test'):
     device = torch.device("cuda" if cfg.use_cuda else "cpu")
-    model_root = cfg.model_path_pretrained_G
-    model_path = os.path.join(model_root, 'G_' + str(cfg.model_step_pretrained_G) + '.pkl')
+    model_path = model_path
     net = torch.load(model_path, map_location=device)
 
     image_color = modelEvalution_out_big(net,
@@ -298,6 +297,11 @@ def out_test(cfg, output_dir='', evaluate_metrics=False, img_name='out_test'):
 def segment_by_out_test(image,model_name):
     print("✅ 传到后端的模型名：", model_name)
 
+    model_path = hf_hub_download(
+        repo_id="weidai00/RIP-AV-sulab",  # 模型库的名字
+        filename=f"G_{model_name}.pkl",  # 文件名
+        repo_type="model"  # 模型库必须写 repo_type
+    )
     cfg.set_dataset(model_name)
     if image is None:
         raise gr.Error("请上传一张图像(upload a fundus image)。")
@@ -306,7 +310,7 @@ def segment_by_out_test(image,model_name):
     temp_path = f"./examples/tmp_upload_{timestamp}.png"
     image.save(temp_path)
 
-    image_color = out_test(cfg, output_dir='', evaluate_metrics=False, img_name=temp_path)
+    image_color = out_test(cfg,model_path=model_path, output_dir='', evaluate_metrics=False, img_name=temp_path)
     return Image.fromarray(image_color)
 
 def gradio_interface():
@@ -361,7 +365,7 @@ def gradio_interface():
             label="示例图像",
             examples_per_page=5
         )
-        with gr.Accordion("📖 模型说明desciption（点击展开）", open=False):
+        with gr.Accordion("📖 模型说明-Description（点击展开）", open=False):
             gr.Markdown(model_info_md)
 
         # 功能连接