Commit 72b87b4e authored by zhengyaoqiu's avatar zhengyaoqiu

优化

parent 0bafc252
*.pth
*.png
*.pt
\ No newline at end of file
......@@ -20,13 +20,14 @@ class FeatureExtractor:
device = "cpu"
self.model, self.preprocess = self.init_model(device, model_name)
self.device = device
model_path = Path(__file__).parent.absolute() / "../models/schp/checkpoints/exp-schp-201908261155-lip.pth"
model_path = Path(__file__).parent.absolute() / "../models/checkpoints/exp-schp-201908261155-lip.pth"
self.schp_model = HumanParsingModel(model_path, dataset='lip')
@staticmethod
def init_model(device="xpu" if torch.xpu.is_available() else "cpu", model_name="ViT-B/32"):
print(f"创建并初始化 CLIP 模型: {model_name} 在设备: {device}")
model, preprocess = clip.load(model_name, device=device)
download_root = Path(__file__).parent.absolute() / "../models"
model, preprocess = clip.load(model_name, device=device, download_root=str(download_root))
return model, preprocess
@staticmethod
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment