Spaces:
Build error
Build error
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| from insightface.app import FaceAnalysis | |
| from huggingface_hub import hf_hub_download | |
| from utils.bisenet import BiSeNet | |
| class FacePreprocessor: | |
| def __init__(self, device='cuda'): | |
| self.device = device if torch.cuda.is_available() else 'cpu' | |
| print(f"🔧 Initializing Face Preprocessor on: {self.device}") | |
| # Initialize InsightFace with better error handling | |
| print(" - Loading InsightFace...") | |
| try: | |
| self.app = FaceAnalysis( | |
| name='antelopev2', | |
| root='./third_party_files', | |
| providers=['CPUExecutionProvider'] | |
| ) | |
| self.app.prepare(ctx_id=0, det_size=(640, 640)) | |
| print(" ✅ InsightFace loaded successfully") | |
| except AssertionError as e: | |
| print(" ⚠️ InsightFace antelopev2 failed, trying buffalo_l...") | |
| try: | |
| # Fallback to buffalo_l model | |
| self.app = FaceAnalysis( | |
| name='buffalo_l', | |
| root='./third_party_files', | |
| providers=['CPUExecutionProvider'] | |
| ) | |
| self.app.prepare(ctx_id=0, det_size=(640, 640)) | |
| print(" ✅ InsightFace buffalo_l loaded successfully") | |
| except Exception as e2: | |
| print(f" ❌ Both InsightFace models failed: {str(e2)}") | |
| raise RuntimeError("Could not load InsightFace models") from e2 | |
| # Initialize BiSeNet for segmentation | |
| print(" - Loading BiSeNet for background extraction...") | |
| self.bisenet = BiSeNet(n_classes=19).to(self.device) | |
| # Download BiSeNet weights from YOUR affecto-faceparsing repo | |
| print(" - Downloading BiSeNet weights from HuggingFace...") | |
| try: | |
| bisenet_path = hf_hub_download( | |
| repo_id="gauravvjhaa/affecto-faceparsing", # ← YOUR REPO | |
| filename="79999_iter.pth", | |
| cache_dir="./models" | |
| ) | |
| print(f" - Loading BiSeNet from: {bisenet_path}") | |
| self.bisenet.load_state_dict( | |
| torch.load(bisenet_path, map_location=self.device) | |
| ) | |
| print(" ✅ BiSeNet loaded successfully!") | |
| except Exception as e: | |
| print(f" ⚠️ BiSeNet download error: {str(e)}") | |
| print(f" ⚠️ Error details: {repr(e)}") | |
| print(" ⚠️ Using random initialization (background extraction will be poor)") | |
| self.bisenet.eval() | |
| # Transforms | |
| self.to_tensor = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) | |
| ]) | |
| print("✅ Preprocessor ready!") | |
| def detect_and_crop_face(self, image): | |
| """Detect face and crop to 512x512""" | |
| print(" 📸 Detecting face...") | |
| img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| try: | |
| faces = self.app.get(img_cv) | |
| except Exception as e: | |
| print(f" ⚠️ Face detection error: {str(e)}") | |
| raise ValueError(f"Face detection failed: {str(e)}") | |
| if len(faces) == 0: | |
| raise ValueError("❌ No face detected!") | |
| if len(faces) > 1: | |
| print(" ⚠️ Multiple faces, using largest") | |
| face = sorted(faces, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1] | |
| bbox = face['bbox'] | |
| x1, y1, x2, y2 = bbox | |
| w = x2 - x1 | |
| h = y2 - y1 | |
| center_x = (x1 + x2) / 2 | |
| center_y = (y1 + y2) / 2 | |
| size = max(w, h) * 1.5 | |
| x1_new = int(center_x - size / 2) | |
| y1_new = int(center_y - size / 2) | |
| x2_new = int(center_x + size / 2) | |
| y2_new = int(center_y + size / 2) | |
| # Handle boundaries | |
| h_img, w_img = img_cv.shape[:2] | |
| x1_new = max(0, x1_new) | |
| y1_new = max(0, y1_new) | |
| x2_new = min(w_img, x2_new) | |
| y2_new = min(h_img, y2_new) | |
| img_crop = img_cv[y1_new:y2_new, x1_new:x2_new] | |
| # Check if crop is valid | |
| if img_crop.size == 0: | |
| raise ValueError("Invalid crop region") | |
| img_crop_512 = cv2.resize(img_crop, (512, 512), interpolation=cv2.INTER_LINEAR) | |
| img_pil = Image.fromarray(cv2.cvtColor(img_crop_512, cv2.COLOR_BGR2RGB)) | |
| print(" ✅ Face cropped to 512x512") | |
| return img_pil | |
| def extract_background(self, image): | |
| """Extract background using BiSeNet""" | |
| print(" 🎨 Extracting background...") | |
| img_tensor = self.to_tensor(image).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| out = self.bisenet(img_tensor)[0] | |
| parsing = out.squeeze(0).cpu().numpy().argmax(0) | |
| img_np = np.array(image) | |
| mask = np.zeros_like(img_np) | |
| # Set face regions to black (classes 1-18 are face parts) | |
| for i in range(1, 19): | |
| mask[parsing == i] = [0, 0, 0] | |
| # Background is class 0 | |
| mask[parsing == 0] = img_np[parsing == 0] | |
| bg_pil = Image.fromarray(mask) | |
| print(" ✅ Background extracted") | |
| return bg_pil | |
| def preprocess(self, image): | |
| """Full preprocessing pipeline""" | |
| try: | |
| cropped = self.detect_and_crop_face(image) | |
| bg = self.extract_background(cropped) | |
| return cropped, bg | |
| except Exception as e: | |
| print(f"❌ Preprocessing error: {str(e)}") | |
| print(" ⚠️ Using fallback (simple resize)") | |
| img_resized = image.resize((512, 512), Image.LANCZOS) | |
| from PIL import ImageFilter | |
| bg = img_resized.filter(ImageFilter.GaussianBlur(radius=10)) | |
| return img_resized, bg |