104 lines
4.3 KiB
Python
104 lines
4.3 KiB
Python
import os
|
|
import zipfile
|
|
from reactor_utils import download # Твоя функция скачивания из utils
|
|
from .face_objects import Face
|
|
from .face_models import SCRFD, ArcFaceONNX, Attribute, Landmark
|
|
from scripts.reactor_logger import logger
|
|
|
|
class ReActorFaceAnalysis:
|
|
"""
|
|
Главный класс-оркестратор.
|
|
Берет картинку, находит лица, определяет пол/возраст и вычисляет эмбеддинги.
|
|
"""
|
|
def __init__(self, name="buffalo_l", root="./models/insightface", providers=None):
|
|
self.name = name
|
|
self.root = root
|
|
self.providers = providers or ["CPUExecutionProvider"]
|
|
self.models = {}
|
|
|
|
model_dir = os.path.join(root, "models", name)
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
|
|
det_file = os.path.join(model_dir, "det_10g.onnx")
|
|
rec_file = os.path.join(model_dir, "w600k_r50.onnx")
|
|
attr_file = os.path.join(model_dir, "genderage.onnx")
|
|
lmk2d_file = os.path.join(model_dir, "2d106det.onnx")
|
|
lmk3d_file = os.path.join(model_dir, "1k3d68.onnx")
|
|
|
|
# Если файлов нет - качаем архив
|
|
if not (os.path.exists(det_file) and os.path.exists(rec_file) and os.path.exists(attr_file)):
|
|
zip_url = "https://huggingface.co/datasets/Gourieff/ReActor/resolve/main/models/buffalo_l.zip"
|
|
zip_path = os.path.join(model_dir, f"{name}.zip")
|
|
|
|
logger.status(f"Downloading {name} models archive...")
|
|
download(zip_url, zip_path, f"{name}.zip")
|
|
|
|
logger.status(f"Extracting {name} models...")
|
|
try:
|
|
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
|
zip_ref.extractall(model_dir)
|
|
logger.status("Extraction completed!")
|
|
except zipfile.BadZipFile:
|
|
logger.error("Downloaded zip file is corrupted. Please try again.")
|
|
finally:
|
|
# В любом случае пытаемся удалить архив, чтобы не занимать место
|
|
if os.path.exists(zip_path):
|
|
os.remove(zip_path)
|
|
|
|
# Инициализируем только те модели, которые физически есть в папке
|
|
if os.path.exists(det_file):
|
|
self.models["detection"] = SCRFD(det_file, providers=self.providers)
|
|
if os.path.exists(rec_file):
|
|
self.models["recognition"] = ArcFaceONNX(rec_file, providers=self.providers)
|
|
if os.path.exists(attr_file):
|
|
self.models["attribute"] = Attribute(attr_file, providers=self.providers)
|
|
if os.path.exists(lmk2d_file):
|
|
self.models["landmark_2d"] = Landmark(lmk2d_file, providers=self.providers)
|
|
if os.path.exists(lmk3d_file):
|
|
self.models["landmark_3d"] = Landmark(lmk3d_file, providers=self.providers)
|
|
|
|
if "detection" not in self.models:
|
|
raise FileNotFoundError(
|
|
f"Detection model (det_10g.onnx) not found at {det_file}. "
|
|
"Please ensure the buffalo_l models are downloaded and extracted properly."
|
|
)
|
|
|
|
def prepare(self, ctx_id=0, det_size=(640, 640), det_thresh=0.5):
|
|
self.det_size = det_size
|
|
self.det_thresh = det_thresh
|
|
|
|
def get(self, img, max_num=0):
|
|
bboxes, kpss = self.models["detection"].detect(
|
|
img,
|
|
det_thresh=self.det_thresh,
|
|
input_size=self.det_size,
|
|
max_num=max_num
|
|
)
|
|
|
|
if bboxes.shape[0] == 0:
|
|
return []
|
|
|
|
ret = []
|
|
for i in range(bboxes.shape[0]):
|
|
bbox = bboxes[i, 0:4]
|
|
det_score = bboxes[i, 4]
|
|
kps = kpss[i] if kpss is not None else None
|
|
|
|
face = Face(bbox=bbox, kps=kps, det_score=det_score)
|
|
|
|
if "attribute" in self.models:
|
|
self.models["attribute"].get(img, face)
|
|
|
|
if "recognition" in self.models:
|
|
self.models["recognition"].get(img, face)
|
|
|
|
if "landmark_2d" in self.models:
|
|
self.models["landmark_2d"].get(img, face)
|
|
|
|
if "landmark_3d" in self.models:
|
|
self.models["landmark_3d"].get(img, face)
|
|
|
|
ret.append(face)
|
|
|
|
return ret
|