Spaces:
Running on Zero
Running on Zero
| """ | |
| CLIP + face embeddings. Same models/config as your Embedding_expansion | |
| notebook: CLIP ViT-H-14 (laion2b_s32b_b79k, 1024-d, L2-normalized) and | |
| InsightFace buffalo_l (512-d). Lazily loaded once on first use instead | |
| of at notebook-cell/import time, since a Gradio app shouldn't pay this | |
| cost unless Part 2 is actually run. | |
| """ | |
| import numpy as np | |
| import torch | |
| CLIP_MODEL_NAME = "ViT-H-14" | |
| CLIP_PRETRAINED = "laion2b_s32b_b79k" | |
| CLIP_DIM = 1024 | |
| FACE_DIM = 512 | |
| _clip_model = None | |
| _clip_preprocess = None | |
| _device = None | |
| _face_app = None | |
| def _get_device(): | |
| global _device | |
| if _device is None: | |
| _device = "cuda" if torch.cuda.is_available() else "cpu" | |
| return _device | |
| def _get_clip(): | |
| global _clip_model, _clip_preprocess | |
| if _clip_model is None: | |
| import open_clip | |
| device = _get_device() | |
| _clip_model, _, _clip_preprocess = open_clip.create_model_and_transforms( | |
| CLIP_MODEL_NAME, pretrained=CLIP_PRETRAINED | |
| ) | |
| _clip_model = _clip_model.to(device).eval() | |
| return _clip_model, _clip_preprocess | |
| def _get_face_app(): | |
| global _face_app | |
| if _face_app is None: | |
| from insightface.app import FaceAnalysis | |
| _face_app = FaceAnalysis( | |
| name="buffalo_l", | |
| providers=["CUDAExecutionProvider", "CPUExecutionProvider"], | |
| ) | |
| _face_app.prepare(ctx_id=0, det_size=(640, 640)) | |
| return _face_app | |
| def embed_clip(pil_image): | |
| """Returns a (1024,) float32 L2-normalized CLIP embedding.""" | |
| model, preprocess = _get_clip() | |
| x = preprocess(pil_image).unsqueeze(0).to(_get_device()) | |
| f = model.encode_image(x) | |
| f = f / f.norm(dim=-1, keepdim=True) | |
| return f.squeeze(0).cpu().numpy().astype(np.float32) | |
| def embed_face(pil_image): | |
| """ | |
| Returns (emb, det_score, bbox_xywh, n_faces, status). | |
| status: "ok" | "ok_multi" (largest face used) | "no_face" | |
| """ | |
| face_app = _get_face_app() | |
| rgb = np.array(pil_image.convert("RGB")) | |
| bgr = rgb[:, :, ::-1].copy() | |
| faces = face_app.get(bgr) | |
| if not faces: | |
| return None, None, None, 0, "no_face" | |
| def area(f): | |
| x1, y1, x2, y2 = f.bbox | |
| return (x2 - x1) * (y2 - y1) | |
| face = max(faces, key=area) if len(faces) > 1 else faces[0] | |
| emb = face.normed_embedding.astype(np.float32) | |
| x1, y1, x2, y2 = [float(v) for v in face.bbox] | |
| bbox = [x1, y1, x2 - x1, y2 - y1] | |
| status = "ok" if len(faces) == 1 else "ok_multi" | |
| return emb, float(face.det_score), bbox, len(faces), status | |