Update app.py

app.py

@@ -9,6 +9,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 from openai import OpenAI
 from huggingface_hub import hf_hub_download
+from segment_anything import SamPredictor, sam_model_registry
+from yolo_world.models.detectors import build_detector
+from mmcv import Config
 
 # 初始化模型
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
@@ -16,13 +19,29 @@ clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 sam_checkpoint = hf_hub_download(
-    repo_id="facebook/sam", filename="sam_vit_h_4b8939.pth", use_auth_token=False
+    repo_id="facebook/sam-vit-large",  # 仓库 ID
+    filename="model.safetensors",      # 模型文件名
+    use_auth_token=False              # 公共仓库无需身份验证
 )
 sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
 sam_predictor = SamPredictor(sam)
-yolo_model = YOLO("yolov8x.pt")  # 替换为实际 YOLO 模型路径
-wd_processor = AutoProcessor.from_pretrained("SmilingWolf/wd-v1-4-vit-large-tagger")
-wd_model = AutoModelForImageClassification.from_pretrained("SmilingWolf/wd-v1-4-vit-large-tagger")
+# 从 Hugging Face 下载 YOLO-World 权重
+yolo_checkpoint = hf_hub_download(
+    repo_id="stevengrove/YOLO-World",  # Hugging Face 仓库 ID
+    filename="yolo_world_v2_xl_obj365v1_goldg_cc3mlite_pretrain.pth",  # 模型权重文件名
+    use_auth_token=False  # 公共仓库无需身份验证
+)
+# 加载 YOLO-World 配置文件
+yolo_config = Config.fromfile('path/to/yolo_world_config.py')  # 替换为实际配置文件路径
+# 构建 YOLO-World 模型
+yolo_model = build_detector(yolo_config.model)
+# 加载权重到模型
+checkpoint = torch.load(yolo_checkpoint, map_location="cpu")  # 使用 CPU 加载权重，后续可以转移到 GPU
+yolo_model.load_state_dict(checkpoint["state_dict"])
+yolo_model.eval()  # 设置为评估模式
+
+wd_processor = AutoProcessor.from_pretrained("SmilingWolf/wd-vit-tagger-v3")
+wd_model = AutoModelForImageClassification.from_pretrained("SmilingWolf/wd-vit-tagger-v3")
 
 # 自动识别图片类型
 def classify_image_type(image):