nodes/smthemex-ComfyUI_Stable_Makeup
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update Retinaface.py

Browse Source 
fix bug
This commit is contained in:
smthemex
2024-07-12 09:13:59 +08:00
committed by GitHub
parent 8184f7eb71
commit c5f36d9bf5
1 changed files with 20 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
facelib/Retinaface/Retinaface.py
+20 -18
View File
@@ -9,7 +9,8 @@ from .utils.box_utils import decode, decode_landmark, prior_box, nms
from .utils.config import cfg_mnet, cfg_re50
from .models.retinaface import RetinaFace
from ..utils import download_weight
import folder_paths
weigths_current_path = os.path.join(folder_paths.models_dir, "stable_makeup")
class FaceDetector:
@@ -29,26 +30,27 @@ class FaceDetector:
if name == 'mobilenet':
cfg = cfg_mnet
model = RetinaFace(cfg=cfg, phase='test')
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'mobilenet.pth')
weight_path = os.path.join(os.path.dirname(file_name), 'weights/mobilenet.pth')
weight_path="F:/test/ComfyUI/models/stable_makeup/mobilenet0.25_Final.pth"
if os.path.isfile(weight_path) == False:
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
download_weight(link='https://drive.google.com/uc?export=download&id=15zP8BP-5IvWXWZoYTNdvUJUiBqZ1hxu1',
file_name=file_name,
verbose=verbose)
os.rename(file_name, weight_path)
if not os.path.exists( weight_path):
try:
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
download_weight(
link='https://drive.google.com/uc?export=download&id=15zP8BP-5IvWXWZoYTNdvUJUiBqZ1hxu1',
file_name=weight_path,
verbose=verbose)
except MaxRetryError as e:
print("看到这个是因为外网不通,无法下载模型，connect Error ",e)
elif name == 'resnet':
cfg = cfg_re50
model = RetinaFace(cfg=cfg, phase='test')
file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'resnet50.pth')
weight_path = os.path.join(os.path.dirname(file_name), 'weights/resnet50.pth')
if os.path.isfile(weight_path) == False:
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
download_weight(link='https://www.dropbox.com/s/8sxkgc9voel6ost/resnet50.pth?dl=1',
file_name=file_name,
verbose=verbose)
os.rename(file_name, weight_path)
if not os.path.exists( weight_path):
try:
os.makedirs(os.path.split(weight_path)[0], exist_ok=True)
download_weight(link='https://www.dropbox.com/s/8sxkgc9voel6ost/resnet50.pth?dl=1',
file_name=weight_path,
verbose=verbose)
except MaxRetryError as e:
print("看到这个是因为外网不通,无法下载模型，connect Error ",e)
else:
exit('FaceDetector Exit: model name can be either mobilenet or resnet')