Spaces:
Running on Zero
Running on Zero
| import os | |
| import random | |
| from functools import lru_cache | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from segment_anything import SamAutomaticMaskGenerator, sam_model_registry | |
| from ultralytics import YOLO | |
| SIZE = (1024, 1024) | |
| NAIL_REPO, NAIL_FILE = "mnemic/nails_seg_yolov8", "nails_seg_s_yolov8_v1.pt" | |
| SAM_REPO, SAM_FILE = "ybelkada/segment-anything", "checkpoints/sam_vit_b_01ec64.pth" | |
| FLUX_MODEL = "black-forest-labs/FLUX.1-Kontext-dev" | |
| def rgba(x): | |
| if isinstance(x, np.ndarray): | |
| x = Image.fromarray(x) | |
| return x.convert("RGBA") | |
| def as_hand(x): | |
| if isinstance(x, dict): | |
| x = x.get("background") | |
| if x is None: | |
| raise gr.Error("请先上传手部图片。") | |
| return rgba(x).resize(SIZE, Image.Resampling.LANCZOS) | |
| def editor_value(hand, mask=None): | |
| hand = as_hand(hand) | |
| if mask is None: | |
| return {"background": hand, "layers": [], "composite": hand} | |
| layer = np.zeros((SIZE[1], SIZE[0], 4), np.uint8) | |
| layer[..., :3] = (255, 47, 146) | |
| layer[..., 3] = (mask > 0).astype(np.uint8) * 150 | |
| return {"background": hand, "layers": [Image.fromarray(layer)], "composite": hand} | |
| def painted_mask(editor): | |
| if not isinstance(editor, dict) or not editor.get("layers"): | |
| raise gr.Error("请先自动识别甲面，或用画笔涂出甲面。") | |
| result = np.zeros((SIZE[1], SIZE[0]), np.uint8) | |
| for layer in editor["layers"]: | |
| if layer is not None: | |
| alpha = np.asarray(rgba(layer).resize(SIZE, Image.Resampling.NEAREST))[..., 3] | |
| result = np.maximum(result, alpha) | |
| result = (result > 16).astype(np.uint8) * 255 | |
| return cv2.morphologyEx(result, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) | |
| def components(mask, minimum): | |
| n, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8) | |
| out = [] | |
| for i in range(1, n): | |
| x, y, w, h, area = stats[i] | |
| if area >= int(minimum): | |
| out.append((np.where(labels == i, 255, 0).astype(np.uint8), int(area), (x, y, w, h))) | |
| return sorted(out, key=lambda item: item[1], reverse=True)[:5] | |
| def quad(mask): | |
| contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| points = cv2.boxPoints(cv2.minAreaRect(max(contours, key=cv2.contourArea))).astype(np.float32) | |
| center = points.mean(0) | |
| points = points[np.argsort(np.arctan2(points[:, 1] - center[1], points[:, 0] - center[0]))] | |
| return np.roll(points, -np.argmin(points[:, 0] + points[:, 1]), 0) | |
| # ---------------- Automatic nail detection ---------------- | |
| def get_nail_model(): | |
| return YOLO(hf_hub_download(NAIL_REPO, NAIL_FILE)) | |
| def detect_nails(editor): | |
| hand = as_hand(editor) | |
| pred = get_nail_model().predict(np.asarray(hand.convert("RGB")), imgsz=1024, conf=0.18, retina_masks=True, device=0, verbose=False)[0] | |
| mask = np.zeros((SIZE[1], SIZE[0]), np.uint8) | |
| if pred.masks is not None: | |
| for item in pred.masks.data.detach().float().cpu().numpy(): | |
| item = cv2.resize(item, SIZE, interpolation=cv2.INTER_NEAREST) | |
| mask = np.maximum(mask, (item > 0.5).astype(np.uint8) * 255) | |
| if not components(mask, 120): | |
| raise gr.Error("未识别到甲面。请使用五指清晰、无遮挡的单手图片，或手动用画笔涂甲面。") | |
| return editor_value(hand, mask), Image.fromarray(mask), f"已自动识别 {len(components(mask, 120))} 个甲面候选；可继续用画笔/橡皮擦修正。" | |
| # ---------------- SAM sticker extraction ---------------- | |
| def get_sam(): | |
| model = sam_model_registry["vit_b"](checkpoint=hf_hub_download(SAM_REPO, SAM_FILE)) | |
| model.to("cuda") | |
| model.eval() | |
| return model | |
| def iou(a, b): | |
| union = np.logical_or(a, b).sum() | |
| return 0 if union == 0 else np.logical_and(a, b).sum() / union | |
| def crop_alpha(rgb, segmentation, bbox): | |
| x, y, w, h = map(int, bbox) | |
| p = max(6, round(max(w, h) * .08)) | |
| x0, y0 = max(0, x-p), max(0, y-p) | |
| x1, y1 = min(rgb.shape[1], x+w+p), min(rgb.shape[0], y+h+p) | |
| return Image.fromarray(np.dstack((rgb[y0:y1, x0:x1], segmentation[y0:y1, x0:x1].astype(np.uint8)*255))) | |
| def thumb(image, side=130): | |
| image = image.copy(); image.thumbnail((side, side), Image.Resampling.LANCZOS) | |
| bg = Image.new("RGBA", (side, side), (246, 246, 246, 255)) | |
| bg.alpha_composite(image, ((side-image.width)//2, (side-image.height)//2)) | |
| return bg.convert("RGB") | |
| def extract_stickers(files, min_area, precision): | |
| if not files: | |
| raise gr.Error("请上传至少一张贴纸排版板。") | |
| generator = SamAutomaticMaskGenerator(model=get_sam(), points_per_side=int(precision), pred_iou_thresh=.82, stability_score_thresh=.88, crop_n_layers=1, crop_n_points_downscale_factor=2, min_mask_region_area=max(20, int(min_area)//3)) | |
| pool, gallery, report = [], [], [] | |
| for board_id, path in enumerate(files[:5]): | |
| image = Image.open(path).convert("RGB") | |
| if max(image.size) > 1200: | |
| r = 1200/max(image.size); image = image.resize((round(image.width*r), round(image.height*r)), Image.Resampling.LANCZOS) | |
| rgb = np.asarray(image); total = rgb.shape[0]*rgb.shape[1] | |
| masks = generator.generate(rgb) | |
| masks.sort(key=lambda m: m["area"]*m["predicted_iou"]*m["stability_score"], reverse=True) | |
| kept, before = [], len(pool) | |
| for item in masks: | |
| segment, bbox, area = item["segmentation"], item["bbox"], int(item["area"]) | |
| x, y, w, h = bbox | |
| aspect = max(w/max(h, 1), h/max(w, 1)) | |
| if area < int(min_area) or area > total*.18 or w < 10 or h < 10 or aspect > 9: | |
| continue | |
| if any(iou(segment, old) > .72 for old in kept): | |
| continue | |
| kept.append(segment) | |
| asset = crop_alpha(rgb, segment, bbox) | |
| pool.append({"image": asset, "board": board_id, "area": area}) | |
| if len(gallery) < 80: | |
| gallery.append((thumb(asset), f"板 {board_id+1}")) | |
| if len(pool) >= 160: | |
| break | |
| report.append(f"板{board_id+1}: {len(pool)-before}个") | |
| if not pool: | |
| raise gr.Error("没有可用图案。尝试降低最小图案面积或提高 SAM 分割精度。") | |
| return pool, gallery, "；".join(report) | |
| # ---------------- Deterministic coverage layout ---------------- | |
| def color(value): | |
| try: | |
| value = value.strip().lstrip("#") | |
| return tuple(int(value[i:i+2], 16) for i in (0, 2, 4)) | |
| except Exception: | |
| return (246, 240, 234) | |
| def covered_assets(pool, count, rng): | |
| by_board = {} | |
| for item in pool: | |
| by_board.setdefault(item["board"], []).append(item) | |
| boards = list(by_board); rng.shuffle(boards) | |
| output = [rng.choice(by_board[b]) for b in boards[:count]] | |
| while len(output) < count: | |
| output.append(rng.choice(pool)) | |
| rng.shuffle(output) | |
| return output, len(boards[:count]) | |
| def nail_design(w, h, asset, index, base): | |
| result = Image.new("RGBA", (max(2, w), max(2, h)), (*base, 248)) | |
| item = rgba(asset) | |
| layouts = [(.50, .50, .55, 0), (.50, .37, .45, 0), (.50, .62, .45, 0), (.42, .52, .40, -10), (.58, .52, .40, 10)] | |
| cx, cy, scale, angle = layouts[index % len(layouts)] | |
| factor = max(w, h)*scale/max(item.size) | |
| item = item.resize((max(2, round(item.width*factor)), max(2, round(item.height*factor))), Image.Resampling.LANCZOS) | |
| item = item.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC) | |
| result.alpha_composite(item, (round(w*cx-item.width/2), round(h*cy-item.height/2))) | |
| return result | |
| def apply(canvas, nail_mask, design, opacity, gloss): | |
| H, W = nail_mask.shape; x, y, w, h = cv2.boundingRect(nail_mask); p=max(8,round(max(w,h)*.08)) | |
| x0,y0,x1,y1=max(0,x-p),max(0,y-p),min(W,x+w+p),min(H,y+h+p) | |
| source=np.asarray(rgba(design).resize((max(2,w),max(2,h)),Image.Resampling.LANCZOS)); sh,sw=source.shape[:2] | |
| target=quad(nail_mask)-np.array([x0,y0],np.float32) | |
| M=cv2.getPerspectiveTransform(np.float32([[0,0],[sw-1,0],[sw-1,sh-1],[0,sh-1]]),target) | |
| warped=cv2.warpPerspective(source,M,(x1-x0,y1-y0),flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT) | |
| local=nail_mask[y0:y1,x0:x1] | |
| alpha=(warped[...,3].astype(np.float32)/255*cv2.GaussianBlur(local,(0,0),1.2)/255*float(opacity))[...,None] | |
| dst=canvas[y0:y1,x0:x1].astype(np.float32); canvas[y0:y1,x0:x1]=(dst*(1-alpha)+warped[...,:3]*alpha).astype(np.uint8) | |
| if gloss: | |
| yy,xx=np.mgrid[:local.shape[0],:local.shape[1]] | |
| shine=np.exp(-(((xx-local.shape[1]*.42)/max(4,local.shape[1]*.18))**2+((yy-local.shape[0]*.25)/max(4,local.shape[0]*.34))**2)*2.2)*local/255*.27*float(gloss) | |
| dst=canvas[y0:y1,x0:x1].astype(np.float32); canvas[y0:y1,x0:x1]=np.clip(dst+shine[...,None]*255,0,255).astype(np.uint8) | |
| def generate(editor, pool, seed, base, opacity, gloss, min_nail): | |
| if not pool: raise gr.Error("请先 AI 提取图案池。") | |
| hand=as_hand(editor); mask=painted_mask(editor); nails=components(mask,min_nail) | |
| if not nails: raise gr.Error("没有甲面。请自动识别或手动画出甲面。") | |
| actual=random.SystemRandom().randint(1,2147483647) if int(seed or 0)==0 else int(seed) | |
| chosen, covered=covered_assets(pool,len(nails),random.Random(actual)) | |
| output=np.asarray(hand.convert("RGB")).copy(); base=color(base) | |
| for i,((nail,_,(_,_,w,h)),item) in enumerate(zip(nails,chosen)): | |
| apply(output,nail,nail_design(w,h,item["image"],i,base),opacity,gloss) | |
| return Image.fromarray(output),Image.fromarray(mask),f"已生成 {len(nails)} 片甲面；{covered} 张素材板均至少出现一次。种子：{actual}",actual | |
| # ---------------- Optional FLUX once, then restore deterministic nails ---------------- | |
| def flux(): | |
| from diffusers import FluxKontextPipeline | |
| token=os.getenv("HF_TOKEN") | |
| if not token: raise gr.Error("FLUX 需要先在 Space Secrets 添加 HF_TOKEN，并在模型页接受许可。") | |
| pipe=FluxKontextPipeline.from_pretrained(FLUX_MODEL,torch_dtype=torch.bfloat16,token=token) | |
| pipe.enable_model_cpu_offload() | |
| return pipe | |
| def polish(design, mask_image, enabled, prompt): | |
| if design is None or mask_image is None: raise gr.Error("请先生成保真结果。") | |
| source=rgba(design).resize(SIZE,Image.Resampling.LANCZOS) | |
| if not enabled: return source,"未启用 FLUX；当前为保真试戴图。" | |
| prompt=prompt or "Premium macro nail photography, soft studio lighting, natural skin texture, editorial background. Preserve one hand, five fingers, pose and rings." | |
| with torch.inference_mode(): | |
| image=flux()(image=source.convert("RGB"),prompt=prompt,num_inference_steps=20,guidance_scale=2.5).images[0].convert("RGB") | |
| protected=np.asarray(mask_image.convert("L").resize(SIZE,Image.Resampling.NEAREST))>0 | |
| final=np.asarray(image).copy(); original=np.asarray(source.convert("RGB")); protected=cv2.dilate(protected.astype(np.uint8),np.ones((3,3),np.uint8),iterations=1)>0 | |
| final[protected]=original[protected] | |
| return Image.fromarray(final),"FLUX 已完成一次全图美化；保真甲面已重新覆盖。" | |
| with gr.Blocks(theme=gr.themes.Soft(),title="Nail Sticker Try-On") as demo: | |
| gr.Markdown("# Nail Sticker Try-On · ZeroGPU\n自动识别甲面 → SAM 提取多张贴纸板 → 每张板至少出现在一根指甲 → 保真贴图 → 可选 FLUX 全图美化并回贴甲面。") | |
| pool=gr.State([]) | |
| with gr.Row(): | |
| with gr.Column(): | |
| editor=gr.ImageEditor(label="1. 手图与甲面修正",type="pil",image_mode="RGBA",canvas_size=SIZE,height=720,width=720,transforms=None,brush=gr.Brush(default_size=24,colors=["#ff2f92"],color_mode="fixed"),eraser=gr.Eraser(default_size=24),layers=False) | |
| detect_button=gr.Button("AI 自动识别甲面（ZeroGPU）") | |
| auto_mask=gr.Image(label="自动识别 mask",type="pil",height=200) | |
| detect_status=gr.Textbox(label="甲面识别状态",interactive=False) | |
| with gr.Column(): | |
| boards=gr.File(label="2. 上传贴纸排版板（最多 5 张）",file_count="multiple",file_types=["image"],type="filepath") | |
| element_area=gr.Slider(50,3000,120,step=10,label="最小图案面积") | |
| precision=gr.Slider(16,32,24,step=4,label="SAM 分割精度") | |
| extract=gr.Button("AI 提取图案池（ZeroGPU）") | |
| extract_status=gr.Textbox(label="图案提取状态",interactive=False) | |
| base=gr.Textbox(value="#F6F0EA",label="甲面底色 HEX") | |
| seed=gr.Number(value=0,precision=0,label="随机种子（0 = 换一组）") | |
| opacity=gr.Slider(.3,1,.92,step=.01,label="贴纸/底色不透明度") | |
| gloss=gr.Slider(0,1,.28,step=.01,label="甲面高光") | |
| min_nail=gr.Slider(80,3000,300,step=20,label="最小甲面面积") | |
| generate_button=gr.Button("生成保真试戴图",variant="primary",size="lg") | |
| gallery=gr.Gallery(label="SAM 候选元素（最多 80 个）",columns=8,height=290) | |
| with gr.Row(): | |
| design=gr.Image(label="保真试戴结果",type="pil",height=650) | |
| final_mask=gr.Image(label="最终甲面 mask",type="pil",height=650) | |
| design_status=gr.Textbox(label="保真生成状态",interactive=False) | |
| gr.Markdown("## 可选：FLUX 全图美化（非商业测试）\n只美化一次背景、手部与光线；代码随后强制覆盖回保真甲面。") | |
| enable=gr.Checkbox(value=False,label="启用 FLUX 全图美化") | |
| prompt=gr.Textbox(value="Premium macro nail photography, soft studio lighting, natural skin texture, elegant editorial background, shallow depth of field. Preserve one hand, five fingers, pose and rings.",label="FLUX 提示词") | |
| polish_button=gr.Button("AI 美化全图并保护甲面（ZeroGPU）") | |
| final=gr.Image(label="最终结果",type="pil",height=700) | |
| polish_status=gr.Textbox(label="美化状态",interactive=False) | |
| # detect.click(detect,[editor],[editor,auto_mask,detect_status]) | |
| detect_button.click( | |
| detect_nails, | |
| [editor], | |
| [editor, auto_mask, detect_status], | |
| ) | |
| extract.click(extract_stickers,[boards,element_area,precision],[pool,gallery,extract_status]) | |
| generate_button.click(generate,[editor,pool,seed,base,opacity,gloss,min_nail],[design,final_mask,design_status,seed]) | |
| polish_button.click(polish,[design,final_mask,enable,prompt],[final,polish_status]) | |
| if __name__=="__main__": | |
| demo.launch() | |