Spaces:
Running on Zero
Running on Zero
File size: 2,526 Bytes
e967c28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """
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
@torch.no_grad()
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
|