Spaces:
Running
Running
| import sys | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from torchvision import transforms | |
| # 1. Injeção da raiz no sys.path para garantir que os imports funcionam | |
| ROOT_DIR = Path(__file__).resolve().parent.parent | |
| if str(ROOT_DIR) not in sys.path: | |
| sys.path.insert(0, str(ROOT_DIR)) | |
| from models.bisnet import BiSeNet | |
| class FaceSegmenter: | |
| def __init__(self, model_path, device="cuda"): | |
| self.device = device | |
| self.n_classes = 19 | |
| self.model = BiSeNet(n_classes=self.n_classes) | |
| # Carregamento robusto dos pesos | |
| state_dict = torch.load(model_path, map_location=device) | |
| self.model.load_state_dict(state_dict) | |
| self.model.to(self.device).eval() | |
| self.transform = transforms.Compose( | |
| [ | |
| transforms.Resize((512, 512)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), | |
| ] | |
| ) | |
| def get_masks(self, img_rgb): | |
| h, w = img_rgb.shape[:2] | |
| img_tensor = ( | |
| self.transform(Image.fromarray(img_rgb)).unsqueeze(0).to(self.device) | |
| ) | |
| with torch.no_grad(): | |
| output = self.model(img_tensor)[0] | |
| # O output do BiSeNet é [1, 19, 512, 512] | |
| parsing = output.squeeze(0).cpu().numpy().argmax(0) | |
| # Redimensionar a segmentação para o tamanho original da imagem | |
| parsing = cv2.resize( | |
| parsing.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST | |
| ) | |
| masks = { | |
| "sobrancelha_esq": (parsing == 2).astype(np.float32), | |
| "sobrancelha_dir": (parsing == 3).astype(np.float32), | |
| "olho_esq": (parsing == 4).astype(np.float32), | |
| "olho_dir": (parsing == 5).astype(np.float32), | |
| "orelha_esq": (parsing == 7).astype(np.float32), | |
| "orelha_dir": (parsing == 8).astype(np.float32), | |
| "nariz": (parsing == 10).astype(np.float32), | |
| "boca": ((parsing == 11) | (parsing == 12) | (parsing == 13)).astype( | |
| np.float32 | |
| ), | |
| "pescoco": (parsing == 14).astype(np.float32), | |
| "cabelo": (parsing == 17).astype(np.float32), | |
| "pele": (parsing == 1).astype(np.float32), | |
| } | |
| # ── Lógica de exclusão e zonas dinâmicas ── | |
| free_skin = (parsing == 1).astype(np.uint8) | |
| # Testa: acima das sobrancelhas | |
| brows = ((parsing == 2) | (parsing == 3)).astype(np.uint8) | |
| if brows.any(): | |
| brow_top_y = int(np.where(brows)[0].min()) | |
| testa = np.zeros((h, w), dtype=np.uint8) | |
| testa[:brow_top_y, :] = free_skin[:brow_top_y, :] | |
| masks["testa"] = testa.astype(np.float32) | |
| else: | |
| masks["testa"] = np.zeros((h, w), dtype=np.float32) | |
| # Bochechas: exclusão mútua usando coordenadas | |
| eye_y = max( | |
| np.where(parsing == 4)[0].max() if (parsing == 4).any() else 0, | |
| np.where(parsing == 5)[0].max() if (parsing == 5).any() else 0, | |
| ) | |
| mouth_top = int( | |
| np.where((parsing >= 11) & (parsing <= 13))[0].min() | |
| if ((parsing >= 11) & (parsing <= 13)).any() | |
| else h | |
| ) | |
| nose_min_x = int( | |
| np.where(parsing == 10)[1].min() if (parsing == 10).any() else w // 2 | |
| ) | |
| nose_max_x = int( | |
| np.where(parsing == 10)[1].max() if (parsing == 10).any() else w // 2 | |
| ) | |
| y_coords, x_coords = np.indices((h, w)) | |
| # Bochecha Esquerda | |
| mask_cheek_l = ( | |
| (free_skin == 1) | |
| & (y_coords > eye_y) | |
| & (y_coords < mouth_top) | |
| & (x_coords < nose_min_x) | |
| ) | |
| # Erosão para não invadir o nariz/olhos | |
| kernel = np.ones((3, 3), np.uint8) | |
| mask_cheek_l = cv2.erode(mask_cheek_l.astype(np.uint8), kernel, iterations=2) | |
| masks["bochecha_esq"] = mask_cheek_l.astype(np.float32) | |
| # Bochecha Direita | |
| mask_cheek_r = ( | |
| (free_skin == 1) | |
| & (y_coords > eye_y) | |
| & (y_coords < mouth_top) | |
| & (x_coords > nose_max_x) | |
| ) | |
| mask_cheek_r = cv2.erode(mask_cheek_r.astype(np.uint8), kernel, iterations=2) | |
| masks["bochecha_dir"] = mask_cheek_r.astype(np.float32) | |
| return masks | |