CoolFace
Apppublic

feiv4020/silhouette-api

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
avatar_engine.py406 linesDownload Raw Back to root
1# ─── SILHOUETTE AVATAR ENGINE ─────────────────────────────────────────────────2# Converts a single user photo into a parametric 3D body mesh3 4import cv25import numpy as np6import mediapipe as mp7import trimesh8from PIL import Image9from rembg import remove10import base6411import io12import json13from dataclasses import dataclass14from typing import Tuple, Optional, List15 16mp_pose     = mp.solutions.pose17mp_face     = mp.solutions.face_detection18mp_segment  = mp.solutions.selfie_segmentation19 20 21# ─── DATA STRUCTURES ──────────────────────────────────────────────────────────22 23@dataclass24class BodyMeasurements:25    height_px:        float26    shoulder_width:   float27    chest_width:      float28    waist_width:      float29    hip_width:        float30    inseam_length:    float31    arm_length:       float32    neck_width:       float33    skin_tone:        Tuple[int, int, int]   # RGB34    # Normalised ratios (0.0–1.0) for mesh shaping35    shoulder_ratio:   float36    waist_ratio:      float37    hip_ratio:        float38    chest_ratio:      float39 40 41@dataclass42class AvatarMesh:43    vertices:   np.ndarray      # (N, 3) float3244    faces:      np.ndarray      # (F, 3) int3245    uvs:        np.ndarray      # (N, 2) float3246    normals:    np.ndarray      # (N, 3) float3247    skin_tone:  Tuple[int, int, int]48    measurements: BodyMeasurements49 50 51# ─── STEP 1: IMAGE PREPROCESSING ─────────────────────────────────────────────52 53class ImageProcessor:54 55    def __init__(self):56        self.pose      = mp_pose.Pose(57                            static_image_mode=True,58                            model_complexity=2,59                            enable_segmentation=True)60        self.face      = mp_face.FaceDetection(min_detection_confidence=0.5)61        self.segmenter = mp_segment.SelfieSegmentation(model_selection=1)62 63    def load_and_preprocess(self, image_bytes: bytes) -> np.ndarray:64        nparr  = np.frombuffer(image_bytes, np.uint8)65        image  = cv2.imdecode(nparr, cv2.IMREAD_COLOR)66        image  = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)67        # Normalise to 1024px tall for consistent landmark scaling68        h, w   = image.shape[:2]69        scale  = 1024 / h70        image  = cv2.resize(image, (int(w * scale), 1024))71        return image72 73    def remove_background(self, image: np.ndarray) -> np.ndarray:74        pil_img    = Image.fromarray(image)75        removed    = remove(pil_img)      # rembg76        return np.array(removed)77 78    def extract_pose_landmarks(79        self, image: np.ndarray80    ) -> Optional[mp_pose.PoseLandmark]:81        results = self.pose.process(image)82        if not results.pose_landmarks:83            raise ValueError("No human body detected in image.")84        return results.pose_landmarks85 86    def extract_skin_tone(87        self, image: np.ndarray, landmarks88    ) -> Tuple[int, int, int]:89        h, w  = image.shape[:2]90        # Sample from face/neck region for accurate skin tone91        nose  = landmarks.landmark[mp_pose.PoseLandmark.NOSE]92        nx, ny = int(nose.x * w), int(nose.y * h)93        # 20x20 sample around nose94        region = image[95            max(0, ny-10):min(h, ny+10),96            max(0, nx-10):min(w, nx+10)97        ]98        if region.size == 0:99            return (210, 180, 140)   # fallback neutral100        mean = region.mean(axis=(0, 1)).astype(int)101        return (int(mean[0]), int(mean[1]), int(mean[2]))102 103 104# ─── STEP 2: BODY MEASUREMENT ESTIMATION ─────────────────────────────────────105 106class MeasurementEstimator:107 108    # Anthropometric reference ratios (female, averaged)109    SHOULDER_TO_HEIGHT  = 0.259110    WAIST_TO_HEIGHT     = 0.181111    HIP_TO_HEIGHT       = 0.191112    CHEST_TO_HEIGHT     = 0.200113    INSEAM_TO_HEIGHT    = 0.471114    NECK_TO_SHOULDER    = 0.210115 116    def estimate(117        self,118        landmarks,119        image_shape: Tuple[int, int]120    ) -> BodyMeasurements:121        h, w = image_shape122        lm   = landmarks.landmark123        L    = mp_pose.PoseLandmark124 125        def px(landmark_id):126            pt = lm[landmark_id]127            return np.array([pt.x * w, pt.y * h])128 129        # Key points130        l_shoulder = px(L.LEFT_SHOULDER)131        r_shoulder = px(L.RIGHT_SHOULDER)132        l_hip      = px(L.LEFT_HIP)133        r_hip      = px(L.RIGHT_HIP)134        l_ankle    = px(L.LEFT_ANKLE)135        r_ankle    = px(L.RIGHT_ANKLE)136        l_wrist    = px(L.LEFT_WRIST)137        l_elbow    = px(L.LEFT_ELBOW)138        nose       = px(L.NOSE)139 140        # Raw pixel measurements141        shoulder_w  = np.linalg.norm(l_shoulder - r_shoulder)142        hip_w       = np.linalg.norm(l_hip - r_hip)143        body_top    = nose[1]144        body_bottom = (l_ankle[1] + r_ankle[1]) / 2145        height_px   = body_bottom - body_top146        inseam_px   = body_bottom - (l_hip[1] + r_hip[1]) / 2147        arm_px      = (148            np.linalg.norm(l_shoulder - l_elbow) +149            np.linalg.norm(l_elbow - l_wrist)150        )151        mid_body    = ((l_shoulder + r_shoulder) / 2 + (l_hip + r_hip) / 2) / 2152        # Waist estimated at midpoint between shoulder and hip153        waist_w     = shoulder_w * 0.72   # typical female ratio154        chest_w     = shoulder_w * 0.88155        neck_w      = shoulder_w * self.NECK_TO_SHOULDER156 157        # Normalised ratios for mesh deformation (around female average)158        # Values > 1.0 = wider than average, < 1.0 = narrower159        avg_shoulder = height_px * self.SHOULDER_TO_HEIGHT160        avg_hip      = height_px * self.HIP_TO_HEIGHT161 162        return BodyMeasurements(163            height_px      = height_px,164            shoulder_width = shoulder_w,165            chest_width    = chest_w,166            waist_width    = waist_w,167            hip_width      = hip_w,168            inseam_length  = inseam_px,169            arm_length     = arm_px,170            neck_width     = neck_w,171            skin_tone      = (0, 0, 0),   # filled by processor172            shoulder_ratio = float(shoulder_w / avg_shoulder),173            waist_ratio    = float(waist_w / (height_px * self.WAIST_TO_HEIGHT)),174            hip_ratio      = float(hip_w / avg_hip),175            chest_ratio    = float(chest_w / (height_px * self.CHEST_TO_HEIGHT)),176        )177 178 179# ─── STEP 3: PARAMETRIC BODY MESH GENERATION ─────────────────────────────────180 181class BodyMeshGenerator:182    """183    Builds a female parametric mesh from body measurements.184    Uses stacked elliptical cross-sections (like a proper185    parametric body model, but without SMPL licensing constraints).186    Each body segment is a tapered elliptic cylinder.187    """188 189    SEGMENTS = 32   # smoothness of cross-sections190 191    def generate(self, m: BodyMeasurements) -> AvatarMesh:192        vertices_list = []193        faces_list    = []194        uvs_list      = []195 196        # Normalise everything to unit height (2.0 Three.js units)197        scale = 2.0 / m.height_px198 199        def sw(px): return px * scale   # scale width200        def sh(px): return px * scale   # scale height201 202        shoulder_r = sw(m.shoulder_width) / 2203        chest_r    = sw(m.chest_width)    / 2204        waist_r    = sw(m.waist_width)    / 2205        hip_r      = sw(m.hip_width)      / 2206        neck_r     = sw(m.neck_width)     / 2207        head_r     = neck_r * 1.85208 209        # Y positions (bottom = 0, top = 2.0)210        y_feet     = 0.0211        y_knee     = sh(m.inseam_length * 0.48)212        y_hip      = sh(m.inseam_length)213        y_waist    = y_hip  + sh(m.height_px * 0.08)214        y_chest    = y_waist + sh(m.height_px * 0.13)215        y_shoulder = y_chest + sh(m.height_px * 0.07)216        y_neck_bot = y_shoulder + sh(m.height_px * 0.03)217        y_neck_top = y_neck_bot + sh(m.height_px * 0.05)218        y_head_bot = y_neck_top219        y_head_top = 2.0220 221        # Body segments: list of (y_bot, r_bot_x, r_bot_z, y_top, r_top_x, r_top_z)222        # x-radius = width, z-radius = depth (depth ≈ 0.6× width for female form)223        DZ = 0.62   # depth ratio224 225        torso_segments = [226            # (y_bot, rx_bot, rz_bot, y_top, rx_top, rz_top, label)227            (y_feet,     hip_r*0.28,      hip_r*0.28*DZ,228             y_knee,     hip_r*0.30,      hip_r*0.30*DZ,     "l_calf"),229            (y_feet,     hip_r*0.28,      hip_r*0.28*DZ,230             y_knee,     hip_r*0.30,      hip_r*0.30*DZ,     "r_calf"),231            (y_knee,     hip_r*0.30,      hip_r*0.30*DZ,232             y_hip,      hip_r*0.42,      hip_r*0.42*DZ,     "l_thigh"),233            (y_knee,     hip_r*0.30,      hip_r*0.30*DZ,234             y_hip,      hip_r*0.42,      hip_r*0.42*DZ,     "r_thigh"),235            (y_hip,      hip_r,           hip_r*DZ,236             y_waist,    waist_r,         waist_r*DZ,         "lower_torso"),237            (y_waist,    waist_r,         waist_r*DZ,238             y_chest,    chest_r,         chest_r*DZ,         "mid_torso"),239            (y_chest,    chest_r,         chest_r*DZ,240             y_shoulder, shoulder_r,      shoulder_r*DZ,      "upper_torso"),241            (y_neck_bot, neck_r,          neck_r,242             y_neck_top, neck_r*0.92,     neck_r*0.92,        "neck"),243        ]244 245        vertex_offset = 0246 247        def add_elliptic_cylinder(248            y_bot, rx_b, rz_b,249            y_top, rx_t, rz_t,250            x_offset=0.0251        ):252            nonlocal vertex_offset253            n     = self.SEGMENTS254            verts = []255            uvs   = []256 257            for i in range(n):258                angle     = 2 * np.pi * i / n259                cos_a     = np.cos(angle)260                sin_a     = np.sin(angle)261                # Bottom ring262                verts.append([x_offset + rx_b*cos_a, y_bot, rz_b*sin_a])263                uvs.append(  [i/n,                    0.0])264                # Top ring265                verts.append([x_offset + rx_t*cos_a, y_top, rz_t*sin_a])266                uvs.append(  [i/n,                    1.0])267 268            faces = []269            for i in range(n):270                b0 = vertex_offset + i*2271                b1 = vertex_offset + ((i+1) % n)*2272                t0 = b0 + 1273                t1 = b1 + 1274                faces.append([b0, t0, b1])275                faces.append([b1, t0, t1])276 277            vertices_list.append(np.array(verts,  dtype=np.float32))278            faces_list.append(   np.array(faces,  dtype=np.int32))279            uvs_list.append(     np.array(uvs,    dtype=np.float32))280            vertex_offset += len(verts)281 282        # Torso (centred)283        for seg in torso_segments[4:]:284            add_elliptic_cylinder(seg[0],seg[1],seg[2],seg[3],seg[4],seg[5])285 286        # Legs (offset left/right)287        leg_offset = hip_r * 0.38288        for seg in torso_segments[:2]:289            add_elliptic_cylinder(290                seg[0],seg[1],seg[2],seg[3],seg[4],seg[5],291                x_offset=-leg_offset292            )293            add_elliptic_cylinder(294                seg[0],seg[1],seg[2],seg[3],seg[4],seg[5],295                x_offset=leg_offset296            )297        for seg in torso_segments[2:4]:298            add_elliptic_cylinder(299                seg[0],seg[1],seg[2],seg[3],seg[4],seg[5],300                x_offset=-leg_offset*0.7301            )302            add_elliptic_cylinder(303                seg[0],seg[1],seg[2],seg[3],seg[4],seg[5],304                x_offset=leg_offset*0.7305            )306 307        # Arms308        arm_r_top = shoulder_r * 0.22309        arm_r_bot = shoulder_r * 0.14310        arm_len   = sh(m.arm_length)311        arm_y_top = y_shoulder312        arm_y_bot = y_shoulder - arm_len313 314        add_elliptic_cylinder(315            arm_y_bot, arm_r_bot, arm_r_bot*0.85,316            arm_y_top, arm_r_top, arm_r_top*0.85,317            x_offset=-(shoulder_r + arm_r_top*0.5)318        )319        add_elliptic_cylinder(320            arm_y_bot, arm_r_bot, arm_r_bot*0.85,321            arm_y_top, arm_r_top, arm_r_top*0.85,322            x_offset= (shoulder_r + arm_r_top*0.5)323        )324 325        # Head — sphere approximated via stacked elliptic rings326        head_h    = y_head_top - y_head_bot327        head_segs = 14328        for i in range(head_segs):329            t0  = i / head_segs330            t1  = (i+1) / head_segs331            ang0 = np.pi * t0332            ang1 = np.pi * t1333            r0  = head_r * np.sin(ang0) * 1.0334            r1  = head_r * np.sin(ang1) * 1.0335            rx0 = r0 * 0.88   # slightly narrower face336            rx1 = r1 * 0.88337            add_elliptic_cylinder(338                y_head_bot + t0*head_h, rx0, r0,339                y_head_bot + t1*head_h, rx1, r1,340            )341 342        # Compile343        all_verts   = np.vstack(vertices_list)344        all_faces   = np.vstack(faces_list)345        all_uvs     = np.vstack(uvs_list)346 347        # Compute normals348        mesh        = trimesh.Trimesh(349                        vertices=all_verts,350                        faces=all_faces,351                        process=False352                      )353        mesh.fix_normals()354        normals     = mesh.vertex_normals.astype(np.float32)355 356        return AvatarMesh(357            vertices     = all_verts,358            faces        = all_faces,359            uvs          = all_uvs,360            normals      = normals,361            skin_tone    = m.skin_tone,362            measurements = m363        )364 365    def to_gltf_dict(self, avatar: AvatarMesh) -> dict:366        """Export as GLTF-compatible JSON for Three.js consumption."""367        r, g, b = [c/255.0 for c in avatar.skin_tone]368        return {369            "vertices":   avatar.vertices.tolist(),370            "faces":      avatar.faces.tolist(),371            "uvs":        avatar.uvs.tolist(),372            "normals":    avatar.normals.tolist(),373            "skin_tone":  {"r": r, "g": g, "b": b},374            "measurements": {375                "shoulder_ratio": avatar.measurements.shoulder_ratio,376                "waist_ratio":    avatar.measurements.waist_ratio,377                "hip_ratio":      avatar.measurements.hip_ratio,378                "chest_ratio":    avatar.measurements.chest_ratio,379            }380        }381 382 383# ─── MASTER AVATAR PIPELINE ───────────────────────────────────────────────────384 385class AvatarPipeline:386 387    def __init__(self):388        self.processor   = ImageProcessor()389        self.estimator   = MeasurementEstimator()390        self.generator   = BodyMeshGenerator()391 392    def run(self, image_bytes: bytes) -> dict:393        # 1. Load394        image       = self.processor.load_and_preprocess(image_bytes)395        # 2. Landmarks396        landmarks   = self.processor.extract_pose_landmarks(image)397        # 3. Skin tone398        skin_tone   = self.processor.extract_skin_tone(image, landmarks)399        # 4. Measurements400        measurements         = self.estimator.estimate(401                                landmarks, image.shape[:2])402        measurements.skin_tone = skin_tone403        # 5. Mesh404        avatar      = self.generator.generate(measurements)405        # 6. Export406        return self.generator.to_gltf_dict(avatar)