Spaces:
Running on Zero
Running on Zero
| # spaces MUST be imported first — before torch or any CUDA-related package | |
| import spaces | |
| import time | |
| import os | |
| import sys | |
| import gc | |
| import io | |
| import json | |
| import base64 | |
| import datetime | |
| import gradio as gr | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| # Add paths | |
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.insert(0, BASE_DIR) | |
| sys.path.insert(0, os.path.join(BASE_DIR, "lib")) | |
| import firebase_config | |
| import firestore_service | |
| import gdown | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| import requests | |
| # --------------------------------------------------------------------------- | |
| # GEMINI VISION — worn vs flat check & pant classification | |
| # --------------------------------------------------------------------------- | |
| _GEMINI_PROMPT = """ | |
| You are an expert garment inspector. Inspect this pants/trousers photo. | |
| Determine two things: | |
| 1. "is_worn_by_person": | |
| - Set to true ONLY IF a living human being or model is actively WEARING the pants (a person's legs, body, torso, shoes, or skin are physically inside the pants). | |
| - Set to false IF the pants are placed empty or laid flat on the floor, bed, table, tile, ground, or hanger (even if the pant legs are spread open). | |
| - If there is NO visible human body inside the pants, is_worn_by_person MUST be false. | |
| 2. "pant_type": | |
| - Classify into ONE category: "Jeans", "Formal Pants", "Chinos", "Cargo Pants", "Track Pants / Joggers", "Linen / Casual Pants", "Shorts", or "Other". | |
| Return ONLY JSON: | |
| {"is_worn_by_person": false, "pant_type": "Jeans"} | |
| """ | |
| def check_pant_with_gemini(img_pil: Image.Image): | |
| """ | |
| Checks if a pant is worn by a person, and classifies the pant type using Gemini. | |
| Non-blocking: if GEMINI_API_KEY is not set or request fails, returns (False, "Pant"). | |
| """ | |
| api_key = os.getenv("GEMINI_API_KEY") | |
| if not api_key: | |
| return False, "Pant" | |
| try: | |
| max_dim = 768 | |
| w, h = img_pil.size | |
| if max(w, h) > max_dim: | |
| s = max_dim / max(w, h) | |
| thumb = img_pil.resize((int(w * s), int(h * s)), Image.BICUBIC) | |
| else: | |
| thumb = img_pil | |
| buf = io.BytesIO() | |
| thumb.convert("RGB").save(buf, format="JPEG", quality=80) | |
| img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") | |
| model_name = os.getenv("GEMINI_MODEL", "gemini-2.5-flash") | |
| url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={api_key}" | |
| payload = { | |
| "contents": [{ | |
| "parts": [ | |
| {"text": _GEMINI_PROMPT}, | |
| {"inline_data": {"mime_type": "image/jpeg", "data": img_b64}} | |
| ] | |
| }], | |
| "generationConfig": { | |
| "response_mime_type": "application/json", | |
| "temperature": 0.0 | |
| } | |
| } | |
| resp = requests.post(url, json=payload, timeout=8) | |
| if resp.status_code == 200: | |
| data = resp.json() | |
| candidates = data.get("candidates", []) | |
| if candidates: | |
| part = candidates[0].get("content", {}).get("parts", [{}])[0] | |
| text = part.get("text", "{}") | |
| parsed = json.loads(text) | |
| is_worn = bool(parsed.get("is_worn_by_person", False)) | |
| pant_type = str(parsed.get("pant_type", "Pant")).strip() | |
| print(f"[Gemini] is_worn={is_worn}, pant_type='{pant_type}'", flush=True) | |
| return is_worn, pant_type | |
| else: | |
| print(f"[Gemini] API error: {resp.status_code} {resp.text[:150]}", flush=True) | |
| except Exception as e: | |
| print(f"[Gemini] Error (continuing without blocking): {e}", flush=True) | |
| return False, "Pant" | |
| # Initialize Firebase | |
| try: | |
| firebase_config.initialize_firebase() | |
| except Exception as e: | |
| print(f"Firebase init error: {e}") | |
| # Initialize Cloudinary | |
| import cloudinary | |
| def init_cloudinary(): | |
| if not os.getenv("CLOUDINARY_CLOUD_NAME"): | |
| return False | |
| import cloudinary.uploader | |
| cloudinary.config( | |
| cloud_name=os.getenv("CLOUDINARY_CLOUD_NAME"), | |
| api_key=os.getenv("CLOUDINARY_API_KEY"), | |
| api_secret=os.getenv("CLOUDINARY_API_SECRET"), | |
| secure=True | |
| ) | |
| return True | |
| init_cloudinary() | |
| # Constants | |
| WAIST_L, WAIST_C, WAIST_R = 6, 0, 30 | |
| OUTER_ANKLE_L, OUTER_ANKLE_R = 14, 16 | |
| INNER_ANKLE_L, INNER_ANKLE_R = 9, 21 | |
| INSEAM = 27 | |
| FALLBACK_INNER_ANKLE_L = 108 | |
| # HRNet model paths | |
| CFG_FILE = os.path.join(BASE_DIR, "experiments/deepfashion2/hrnet/w48_384x288_adam_lr1e-3.yaml") | |
| MODEL_FILE = os.path.join(BASE_DIR, "lib/models/pose_hrnet-w48_384x288-deepfashion2_mAP_0.7017.pth") | |
| MODEL_URL = "https://drive.google.com/file/d/1Srj8YJLu00aZm368sCtEIWB8Vkpr7kXA/view?usp=sharing" | |
| def ensure_model_exists(): | |
| if not os.path.exists(MODEL_FILE): | |
| os.makedirs(os.path.dirname(MODEL_FILE), exist_ok=True) | |
| print("Downloading HRNet model...", flush=True) | |
| gdown.download(MODEL_URL, MODEL_FILE, quiet=True) | |
| ensure_model_exists() | |
| # Load HRNet model (CPU — GPU assigned dynamically via @spaces.GPU) | |
| from config import cfg, update_config | |
| import models | |
| import argparse | |
| args = argparse.Namespace(cfg=CFG_FILE, opts=[], modelDir="", logDir="", dataDir="", prevModelDir="") | |
| update_config(cfg, args) | |
| cfg.freeze() | |
| _hrnet = models.pose_hrnet.get_pose_net(cfg, is_train=False) | |
| _hrnet.load_state_dict(torch.load(MODEL_FILE, map_location="cpu"), strict=False) | |
| _hrnet.eval() | |
| # Load SAM3 model (CPU — GPU assigned dynamically via @spaces.GPU) | |
| from huggingface_hub import hf_hub_download | |
| _sam3_model = None | |
| _sam3_processor = None | |
| cumulative_gpu_time = 0.0 | |
| def init_sam3(): | |
| global _sam3_model, _sam3_processor | |
| try: | |
| from sam3.model_builder import build_sam3_image_model | |
| from sam3.model.sam3_image_processor import Sam3Processor | |
| token = os.getenv("HF_TOKEN") | |
| if not token: | |
| print("WARNING: HF_TOKEN environment variable not set. Downloading SAM3 may fail.") | |
| print("Downloading/loading SAM3 model...", flush=True) | |
| model_path = hf_hub_download(repo_id="facebook/sam3", filename="sam3.pt", token=token) | |
| _sam3_model = build_sam3_image_model(checkpoint_path=model_path) | |
| _sam3_model.eval() | |
| _sam3_processor = Sam3Processor(_sam3_model) | |
| print("SAM3 loaded successfully.", flush=True) | |
| except Exception as e: | |
| print(f"Failed to load SAM3: {e}") | |
| init_sam3() | |
| # --------------------------------------------------------------------------- | |
| # SINGLE @spaces.GPU function — runs both ISNet BG removal and HRNet | |
| # --------------------------------------------------------------------------- | |
| def run_pipeline_gpu(img_pil: Image.Image, prompt_text="pants"): | |
| """ | |
| Runs the full GPU pipeline: | |
| 1. SAM3 garment segmentation (uses GPU if available) | |
| 2. HRNet landmark estimation (uses GPU) | |
| Returns: | |
| img_np : numpy HWC image array (for drawing) | |
| binary_mask : numpy 2-D mask array | |
| heatmaps : numpy array of shape (num_joints, H, W) | |
| img_h, img_w: original dimensions used for landmark scaling | |
| gpu_time : time spent in GPU function this run | |
| """ | |
| global _hrnet, _sam3_model, _sam3_processor, cumulative_gpu_time | |
| t_start = time.time() | |
| # -- Step 1: SAM3 garment segmentation -- | |
| t1 = time.time() | |
| if _sam3_model is not None and _sam3_processor is not None: | |
| _sam3_model.to("cuda" if torch.cuda.is_available() else "cpu") | |
| try: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ZeroGPU does not support bfloat16 — use float16 on GPU, float32 on CPU | |
| dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| with torch.inference_mode(), torch.autocast(device_type=device, dtype=dtype): | |
| inference_state = _sam3_processor.set_image(img_pil) | |
| output = _sam3_processor.set_text_prompt(state=inference_state, prompt=prompt_text) | |
| masks = output["masks"] | |
| scores = output["scores"] | |
| if torch.is_tensor(scores): | |
| scores_np = scores.cpu().numpy() | |
| else: | |
| scores_np = np.asarray(scores) | |
| # Filter masks with confidence > 0.5 (as per the colab notebook) | |
| valid_indices = np.where(scores_np > 0.5)[0] | |
| if len(valid_indices) == 0: | |
| # If nothing > 0.5, fallback to all indices to at least output something | |
| valid_indices = np.arange(len(scores_np)) | |
| # Find the largest mask by area among the valid ones | |
| best_idx = valid_indices[0] | |
| max_area = -1 | |
| if torch.is_tensor(masks): | |
| for idx in valid_indices: | |
| area = masks[idx].sum().item() | |
| if area > max_area: | |
| max_area = area | |
| best_idx = idx | |
| else: | |
| for idx in valid_indices: | |
| area = masks[idx].sum() | |
| if area > max_area: | |
| max_area = area | |
| best_idx = idx | |
| best_mask = masks[best_idx] | |
| if torch.is_tensor(best_mask): | |
| best_mask = best_mask.cpu().numpy() | |
| # Ensure shape is (H, W) | |
| if best_mask.ndim == 3 and best_mask.shape[0] == 1: | |
| best_mask = best_mask[0] | |
| if best_mask.shape != (img_pil.height, img_pil.width): | |
| best_mask = cv2.resize(best_mask.astype(np.uint8), (img_pil.width, img_pil.height), interpolation=cv2.INTER_NEAREST) | |
| mask_np = (best_mask > 0).astype(np.uint8) * 255 | |
| except Exception as e: | |
| print(f"Error during SAM3 inference: {e}", flush=True) | |
| mask_np = np.zeros((img_pil.height, img_pil.width), dtype=np.uint8) | |
| else: | |
| print("SAM3 model not initialized, falling back to empty mask.") | |
| mask_np = np.zeros((img_pil.size[1], img_pil.size[0]), dtype=np.uint8) | |
| isnet_time = time.time() - t1 | |
| print(f"SAM3: {isnet_time:.2f}s", flush=True) | |
| # Generate transparent cutout / white background | |
| img_rgb = img_pil.convert("RGB") | |
| white_bg = Image.new("RGB", img_pil.size, "WHITE") | |
| mask_pil = Image.fromarray(mask_np) | |
| white_bg.paste(img_rgb, mask=mask_pil) | |
| # Portrait orientation helpers | |
| def _portrait(arr_img): | |
| h, w = arr_img.shape[:2] | |
| if w > h: | |
| arr_img = cv2.rotate(arr_img, cv2.ROTATE_90_COUNTERCLOCKWISE) | |
| return arr_img | |
| img_np = _portrait(np.array(white_bg)) | |
| binary_mask = _portrait(mask_np) | |
| img_np_orig = _portrait(np.array(img_rgb)) | |
| # -- Step 2: HRNet landmark estimation -- | |
| t2 = time.time() | |
| H, W = img_np.shape[:2] | |
| inp = cv2.resize(img_np, (288, 384)).astype(np.float32) / 255.0 | |
| inp = inp.transpose(2, 0, 1) | |
| inp_tensor = torch.from_numpy(inp).unsqueeze(0) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| _hrnet = _hrnet.to(device) | |
| with torch.inference_mode(): | |
| heatmaps = _hrnet(inp_tensor.to(device)).squeeze(0).cpu().numpy() | |
| hrnet_time = time.time() - t2 | |
| total = time.time() - t_start | |
| cumulative_gpu_time += total | |
| print(f"HRNet: {hrnet_time:.2f}s", flush=True) | |
| print(f"Total GPU function: {total:.2f}s (Cumulative since boot: {cumulative_gpu_time:.2f}s)", flush=True) | |
| return img_np_orig, binary_mask, heatmaps, H, W, total | |
| # --------------------------------------------------------------------------- | |
| # CPU helpers — landmark decoding, refinement, measurement, drawing | |
| # --------------------------------------------------------------------------- | |
| def decode_landmarks(heatmaps, H, W): | |
| num_joints, hm_h, hm_w = heatmaps.shape | |
| pts = {} | |
| for jid in range(num_joints): | |
| y, x = np.unravel_index(np.argmax(heatmaps[jid]), heatmaps[jid].shape) | |
| pts[jid] = (int(x * W / hm_w), int(y * H / hm_h)) | |
| if FALLBACK_INNER_ANKLE_L in pts and pts[FALLBACK_INNER_ANKLE_L][1] > pts[INNER_ANKLE_L][1]: | |
| pts[INNER_ANKLE_L] = pts[FALLBACK_INNER_ANKLE_L] | |
| return pts | |
| def refine_landmarks(pts, binary_mask, max_snap_dist_ratio=0.08): | |
| refined_pts = pts.copy() | |
| contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) | |
| if not contours: | |
| return refined_pts | |
| largest_contour = max(contours, key=cv2.contourArea) | |
| contour_pts = largest_contour.reshape(-1, 2) | |
| H, W = binary_mask.shape[:2] | |
| max_snap_dist = max_snap_dist_ratio * max(H, W) | |
| boundary_ids = [WAIST_L, WAIST_R, OUTER_ANKLE_L, OUTER_ANKLE_R, INNER_ANKLE_L, INNER_ANKLE_R, INSEAM] | |
| for lid in boundary_ids: | |
| if lid not in pts: | |
| continue | |
| pt = np.array(pts[lid], dtype=np.float32) | |
| dists = np.linalg.norm(contour_pts - pt, axis=1) | |
| min_idx = np.argmin(dists) | |
| if dists[min_idx] <= max_snap_dist: | |
| refined_pts[lid] = (int(contour_pts[min_idx][0]), int(contour_pts[min_idx][1])) | |
| return refined_pts | |
| def round_half(x): | |
| return round(x * 2) / 2 | |
| def dist(p, a, b): | |
| x1, y1 = p[a]; x2, y2 = p[b] | |
| return float(np.hypot(x1 - x2, y1 - y2)) | |
| def draw_measurements(img, pts): | |
| draw = img.copy() | |
| for a, b in [ | |
| (WAIST_L, WAIST_C), (WAIST_C, WAIST_R), | |
| (WAIST_L, OUTER_ANKLE_L), (WAIST_R, OUTER_ANKLE_R), | |
| (INNER_ANKLE_L, OUTER_ANKLE_L), (INNER_ANKLE_R, OUTER_ANKLE_R), | |
| (INNER_ANKLE_L, INSEAM), (INNER_ANKLE_R, INSEAM), | |
| ]: | |
| cv2.line(draw, pts[a], pts[b], (60, 120, 180), 2) | |
| return draw | |
| def upload_image_to_cloud(image_obj, folder): | |
| if not os.getenv("CLOUDINARY_CLOUD_NAME"): | |
| return None | |
| try: | |
| import cloudinary.uploader | |
| if isinstance(image_obj, (np.ndarray, np.generic)): | |
| pil_img = Image.fromarray(image_obj) | |
| buf = io.BytesIO() | |
| pil_img.save(buf, format="JPEG", quality=80) | |
| buf.seek(0) | |
| image_obj = buf | |
| return cloudinary.uploader.upload(image_obj, folder=folder, resource_type="image", timeout=120).get("secure_url") | |
| except Exception as e: | |
| print(f"Image upload failed: {e}") | |
| return None | |
| def verify_user_from_token(token): | |
| if not token: | |
| raise ValueError("No token found") | |
| from firebase_admin import auth as firebase_auth | |
| return firebase_auth.verify_id_token(token)['uid'] | |
| # --------------------------------------------------------------------------- | |
| # Main Gradio handler | |
| # --------------------------------------------------------------------------- | |
| def process_fit(name, waist_in, input_image, request: gr.Request): | |
| if not name: | |
| return None, "Please enter your name.", "", "" | |
| if input_image is None: | |
| return None, "Please upload an image.", "", "" | |
| try: | |
| # Auth | |
| uid = None | |
| if request: | |
| token = request.query_params.get("token") | |
| if token: | |
| try: | |
| uid = verify_user_from_token(token) | |
| except Exception as e: | |
| print(f"Auth failed: {e}") | |
| # Preprocess: EXIF transpose + optional resize | |
| img_pil = ImageOps.exif_transpose(input_image) | |
| orig_w, orig_h = img_pil.size | |
| MAX_DIM = 1024 | |
| scale_x = scale_y = 1.0 | |
| resized = False | |
| if max(orig_w, orig_h) > MAX_DIM: | |
| scale = MAX_DIM / max(orig_w, orig_h) | |
| new_w, new_h = int(orig_w * scale), int(orig_h * scale) | |
| scale_x, scale_y = new_w / orig_w, new_h / orig_h | |
| resized = True | |
| img_cv2 = cv2.cvtColor(np.array(img_pil.convert("RGB")), cv2.COLOR_RGB2BGR) | |
| img_pil = Image.fromarray(cv2.cvtColor( | |
| cv2.resize(img_cv2, (new_w, new_h), interpolation=cv2.INTER_AREA), | |
| cv2.COLOR_BGR2RGB | |
| )) | |
| # Check if pant is worn by a person using Gemini Vision | |
| is_worn, pant_type = check_pant_with_gemini(img_pil) | |
| if is_worn: | |
| return ( | |
| None, | |
| "<div style='background:#fef2f2;border:1.5px solid #ef4444;border-radius:16px;padding:18px;color:#991b1b;font-weight:600;font-size:1rem;'>" | |
| "🚫 <b>Image Rejected:</b> A person appears to be wearing this pant.<br><br>" | |
| "Please lay the pant <b>completely flat on a surface or floor</b> and re-upload.</div>", | |
| "", | |
| "Rejected: Pant is worn by a person" | |
| ) | |
| # Run GPU pipeline (SAM3 + HRNet in one @spaces.GPU call) | |
| img_np_orig, binary_mask, heatmaps, H, W, gpu_time = run_pipeline_gpu(img_pil, prompt_text="pants") | |
| # Verify if a pant is actually detected | |
| contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if not contours or np.sum(binary_mask) < 5000: | |
| return None, "No pant found in the image. Please lay the pant flat and make sure it is fully visible in the photo.", "", "" | |
| # Decode + refine landmarks (CPU) | |
| pts = decode_landmarks(heatmaps, H, W) | |
| pts = refine_landmarks(pts, binary_mask) | |
| # Apply black background and pink overlay for the garment mask | |
| PINK = (255, 105, 180) # RGB pink | |
| ALPHA = 0.55 | |
| black_bg = np.zeros_like(img_np_orig) | |
| result_vis = black_bg.copy() | |
| pink_arr = np.zeros_like(img_np_orig) | |
| pink_arr[:] = PINK | |
| # Blend original pixel + pink overlay where mask is True | |
| mask_bool = binary_mask > 128 | |
| blended = (img_np_orig.astype(np.float32) * (1.0 - ALPHA) + pink_arr.astype(np.float32) * ALPHA).astype(np.uint8) | |
| result_vis[mask_bool] = blended[mask_bool] | |
| # Draw pink contour around the mask | |
| contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| cv2.drawContours(result_vis, contours, -1, PINK, thickness=2, lineType=cv2.LINE_AA) | |
| result_img = draw_measurements(result_vis, pts) | |
| # Scale pts back if image was resized | |
| if resized: | |
| for lid in pts: | |
| pts[lid] = (int(pts[lid][0] / scale_x), int(pts[lid][1] / scale_y)) | |
| # Measurements | |
| waist_px = dist(pts, WAIST_L, WAIST_C) + dist(pts, WAIST_C, WAIST_R) | |
| outseam_px = (dist(pts, WAIST_L, OUTER_ANKLE_L) + dist(pts, WAIST_R, OUTER_ANKLE_R)) / 2 | |
| ankle_px = (dist(pts, INNER_ANKLE_L, OUTER_ANKLE_L) + dist(pts, INNER_ANKLE_R, OUTER_ANKLE_R)) / 2 | |
| inseam_px = (dist(pts, INSEAM, INNER_ANKLE_L) + dist(pts, INSEAM, INNER_ANKLE_R)) / 2 | |
| half_waist = (waist_in / 2) - 1 | |
| outseam_val = round_half(outseam_px / waist_px * half_waist) | |
| ankle_val = round_half(ankle_px / waist_px * half_waist) | |
| inseam_val = round_half(inseam_px / waist_px * half_waist) | |
| result_pil = Image.fromarray(result_img) | |
| # Cloud save (if logged in) | |
| sync_status = "Not logged in — scan not saved to cloud profile." | |
| if uid: | |
| image_url = upload_image_to_cloud(result_img, "dashclothing/scans") | |
| if image_url: | |
| try: | |
| firestore_service.save_prediction(uid, { | |
| "name": name, "imageUrl": image_url, | |
| "waist": waist_in, "outseam": outseam_val, | |
| "inseam": inseam_val, "ankle": ankle_val, | |
| "status": "completed" | |
| }) | |
| sync_status = "✅ Saved to your DASH profile!" | |
| except Exception as ex: | |
| sync_status = f"Upload OK but DB sync failed: {ex}" | |
| else: | |
| sync_status = "Cloudinary upload failed." | |
| # Normalized landmarks for AR 3D unprojection | |
| landmarks_norm = { | |
| "waist_left": {"x": float(pts[WAIST_L][0] / orig_w), "y": float(pts[WAIST_L][1] / orig_h)}, | |
| "waist_center": {"x": float(pts[WAIST_C][0] / orig_w), "y": float(pts[WAIST_C][1] / orig_h)}, | |
| "waist_right": {"x": float(pts[WAIST_R][0] / orig_w), "y": float(pts[WAIST_R][1] / orig_h)}, | |
| "inseam": {"x": float(pts[INSEAM][0] / orig_w), "y": float(pts[INSEAM][1] / orig_h)}, | |
| "outer_ankle_left": {"x": float(pts[OUTER_ANKLE_L][0] / orig_w), "y": float(pts[OUTER_ANKLE_L][1] / orig_h)}, | |
| "outer_ankle_right": {"x": float(pts[OUTER_ANKLE_R][0] / orig_w), "y": float(pts[OUTER_ANKLE_R][1] / orig_h)}, | |
| "inner_ankle_left": {"x": float(pts[INNER_ANKLE_L][0] / orig_w), "y": float(pts[INNER_ANKLE_L][1] / orig_h)}, | |
| "inner_ankle_right": {"x": float(pts[INNER_ANKLE_R][0] / orig_w), "y": float(pts[INNER_ANKLE_R][1] / orig_h)}, | |
| } | |
| # postMessage payload for parent iframe | |
| fit_payload = json.dumps({ | |
| "type": "dashFitResult", "name": name, "waist": waist_in, | |
| "outseam": float(outseam_val), "inseam": float(inseam_val), | |
| "ankle": float(ankle_val), "recommendedSize": str(int(waist_in)), | |
| "pantType": pant_type, | |
| "landmarks": landmarks_norm, | |
| "saved_at": datetime.datetime.now(datetime.timezone.utc).isoformat() | |
| }) | |
| post_msg_html = f""" | |
| <script> | |
| (function(){{ | |
| var p={fit_payload}; | |
| try{{window.top.postMessage(p,'*');}}catch(e){{}} | |
| try{{window.parent.postMessage(p,'*');}}catch(e){{}} | |
| console.log('✅ Fit data:', p); | |
| }})(); | |
| </script>""" | |
| results_html = f""" | |
| <div style="background:#fff;border:1px solid #f4f4f5;border-radius:20px;padding:20px;font-family:sans-serif;box-shadow:0 4px 6px rgba(0,0,0,.05);"> | |
| <h3 style="margin-top:0;color:#18181b;">Results for {name}</h3> | |
| <div style="font-size:.8rem;color:#71717a;text-transform:uppercase;font-weight:700;">Recommended Size</div> | |
| <div style="font-size:3rem;font-weight:800;color:#000;margin:8px 0;">{int(waist_in)}</div> | |
| <div style="display:grid;grid-template-columns:1fr 1fr;gap:12px;border-top:1px solid #f4f4f5;padding-top:12px;"> | |
| <div><div style="font-size:.75rem;color:#a1a1aa;">Waist</div><div style="font-size:1.1rem;font-weight:700;">{waist_in} in</div></div> | |
| <div><div style="font-size:.75rem;color:#a1a1aa;">Outseam</div><div style="font-size:1.1rem;font-weight:700;">{outseam_val} in</div></div> | |
| <div><div style="font-size:.75rem;color:#a1a1aa;">Inseam</div><div style="font-size:1.1rem;font-weight:700;">{inseam_val} in</div></div> | |
| <div><div style="font-size:.75rem;color:#a1a1aa;">Ankle</div><div style="font-size:1.1rem;font-weight:700;">{ankle_val} in</div></div> | |
| </div> | |
| <div style="margin-top:15px; font-size:0.75rem; color:#a1a1aa; text-align:center;"> | |
| ⏱️ ZeroGPU Time Used: {gpu_time:.2f}s (Cumulative: {cumulative_gpu_time:.2f}s) | |
| </div> | |
| <div style="margin-top:10px;text-align:center;"> | |
| <span style="display:inline-block;background:#f4f4f5;border-radius:999px;padding:4px 14px;font-size:0.82rem;font-weight:700;color:#27272a;"> | |
| 👖 {pant_type} | |
| </span> | |
| </div> | |
| </div>""" | |
| return result_pil, results_html, post_msg_html, sync_status | |
| except Exception as e: | |
| import traceback | |
| return None, f"<div style='color:red;'>Error: {e}</div>", "", traceback.format_exc() | |
| def process_fit_json(name, waist_in, input_image, token=""): | |
| # Simple wrapper for API calls that returns pure JSON | |
| if not name or input_image is None: | |
| return {"error": "Missing name or image"} | |
| # Mocking a request object with query_params for token | |
| class DummyRequest: | |
| def __init__(self, t): | |
| self.query_params = {"token": t} | |
| # Run the main process | |
| res = process_fit(name, waist_in, input_image, DummyRequest(token) if token else None) | |
| if res[0] is None: | |
| import re | |
| clean_msg = re.sub(r'<[^>]*>', ' ', str(res[1])) | |
| clean_msg = re.sub(r'\s+', ' ', clean_msg).strip() | |
| clean_msg = clean_msg.replace("🚫 ", "") | |
| return {"error": clean_msg or "Detection failed."} | |
| # Convert PIL image to base64 — resize to max 800px to keep response small | |
| result_img = res[0] | |
| max_dim = 800 | |
| if max(result_img.size) > max_dim: | |
| scale = max_dim / max(result_img.size) | |
| result_img = result_img.resize( | |
| (int(result_img.width * scale), int(result_img.height * scale)), | |
| Image.LANCZOS | |
| ) | |
| buf = io.BytesIO() | |
| result_img.save(buf, format="JPEG", quality=72, optimize=True) | |
| buf.seek(0) | |
| result_b64 = base64.b64encode(buf.read()).decode("utf-8") | |
| # Extract the fit_payload from the post_msg_html script | |
| import re | |
| match = re.search(r'var p=({.*?});', res[2], re.DOTALL) | |
| payload = json.loads(match.group(1)) if match else {} | |
| # gpu time from results_html | |
| gpu_time = 0.0 | |
| cum_time = 0.0 | |
| match_time = re.search(r'ZeroGPU Time Used: ([\d\.]+)s \(Cumulative: ([\d\.]+)s\)', res[1]) | |
| if match_time: | |
| gpu_time = float(match_time.group(1)) | |
| cum_time = float(match_time.group(2)) | |
| return { | |
| "result_image_b64": result_b64, | |
| "landmarks": payload.get("landmarks", {}), | |
| "waist": payload.get("waist", waist_in), | |
| "outseam": payload.get("outseam", 0), | |
| "inseam": payload.get("inseam", 0), | |
| "ankle": payload.get("ankle", 0), | |
| "recommended_size": payload.get("recommendedSize", str(waist_in)), | |
| "pant_type": payload.get("pantType", "Pant"), | |
| "gpu_time": gpu_time, | |
| "cumulative_gpu_time": cum_time, | |
| "sync_status": "saved" if "✅" in res[3] else ("not_logged_in" if "Not logged in" in res[3] else res[3]), | |
| "saved_at": payload.get("saved_at", "") | |
| } | |
| def detect_landmarks(input_image): | |
| """ | |
| Dedicated endpoint for AR measurement. | |
| Takes the garment photo and returns normalized 2D landmarks for 3D unprojection. | |
| """ | |
| if input_image is None: | |
| return {"error": "Missing image"} | |
| return process_fit_json("AR_Measure", 32, input_image, token="") | |
| # --------------------------------------------------------------------------- | |
| # Gradio UI | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(theme=gr.themes.Default(primary_hue="emerald", secondary_hue="indigo")) as demo: | |
| gr.HTML(""" | |
| <div style="text-align:center;margin-bottom:20px;"> | |
| <h1 style="font-weight:800;font-size:2rem;margin-bottom:5px;color:#111827;">DASH AI Fit</h1> | |
| <p style="color:#4b5563;font-size:1.1rem;">Find your perfect size in 10 seconds.</p> | |
| </div>""") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| name_input = gr.Textbox(label="Your Name", placeholder="Enter your name") | |
| waist_input = gr.Number(label="Waist Size (inches)", value=32, minimum=24, maximum=60, step=1) | |
| image_input = gr.Image(label="Upload Flat Laid Pant Image", type="pil", sources=["upload", "webcam"]) | |
| token_input = gr.Textbox(visible=False) # For API | |
| analyze_btn = gr.Button("ANALYZE MY FIT", variant="primary") | |
| with gr.Column(scale=1): | |
| image_output = gr.Image(label="AI Detection Result") | |
| results_panel = gr.HTML(label="Results") | |
| status_panel = gr.Textbox(label="Sync Status", interactive=False) | |
| post_msg_panel = gr.HTML() | |
| api_json_out = gr.JSON(visible=False) | |
| analyze_btn.click( | |
| fn=process_fit, | |
| inputs=[name_input, waist_input, image_input], | |
| outputs=[image_output, results_panel, post_msg_panel, status_panel], | |
| api_name="analyze" | |
| ) | |
| # Hidden button to expose pure JSON API | |
| api_btn = gr.Button("API", visible=False) | |
| api_btn.click( | |
| fn=process_fit_json, | |
| inputs=[name_input, waist_input, image_input, token_input], | |
| outputs=api_json_out, | |
| api_name="analyze_json" | |
| ) | |
| # Dedicated AR Landmark API endpoint | |
| ar_api_btn = gr.Button("AR Landmarks API", visible=False) | |
| ar_api_btn.click( | |
| fn=detect_landmarks, | |
| inputs=[image_input], | |
| outputs=api_json_out, | |
| api_name="detect_landmarks" | |
| ) | |
| demo.queue().launch(show_error=True) | |