app.py

"""
Ghost Mannequin Pro
- Removes person (all body parts including legs/feet)
- Inpaints holes where hands covered clothing
- Front+Back 3D composite with size normalization
"""

from flask import Flask, request, jsonify, send_file
from PIL import Image, ImageFilter
import io, numpy as np, traceback, torch
from scipy.ndimage import binary_dilation, binary_erosion, binary_fill_holes, gaussian_filter

app = Flask(__name__)

# ALL person body parts to remove (expanded list)
PERSON_IDS = {
    2,   # Hair
    11,  # Face
    12,  # Left-leg
    13,  # Right-leg
    14,  # Left-arm
    15,  # Right-arm
    # Note: feet/shoes are 9,10 but those are also shoes we want to keep
    # So we detect shoes separately
}

# IDs that are clothing we KEEP
CLOTHING_IDS = {1,4,5,6,7,8,9,10,16,17}

_seg_processor = None
_seg_model = None
_inpaint_pipe = None

def get_seg_model():
    global _seg_processor, _seg_model
    if _seg_model is None:
        from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation
        print("Loading segmentation model...")
        _seg_processor = SegformerImageProcessor.from_pretrained("sayeed99/segformer_b3_clothes")
        _seg_model = AutoModelForSemanticSegmentation.from_pretrained("sayeed99/segformer_b3_clothes")
        _seg_model.eval()
        print("Segmentation model ready!")
    return _seg_processor, _seg_model

def get_inpaint_pipe():
    global _inpaint_pipe
    if _inpaint_pipe is None:
        from diffusers import StableDiffusionInpaintPipeline
        print("Loading inpainting model...")
        _inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-inpainting",
            torch_dtype=torch.float32,
        )
        _inpaint_pipe.enable_attention_slicing()
        print("Inpainting model ready!")
    return _inpaint_pipe

def remove_bg(img):
    from rembg import remove
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "PNG")
    result = remove(buf.getvalue())
    return Image.open(io.BytesIO(result)).convert("RGBA")

def get_segmentation(img_rgb):
    """Returns full label map"""
    processor, model = get_seg_model()
    inputs = processor(images=img_rgb, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    upsampled = torch.nn.functional.interpolate(
        logits, size=img_rgb.size[::-1], mode="bilinear", align_corners=False
    )
    return upsampled.argmax(dim=1)[0].numpy()  # (H, W)

def get_masks(img_rgb):
    """Returns person_mask and clothing_mask as float arrays"""
    pred = get_segmentation(img_rgb)
    H, W = pred.shape

    person_mask = np.zeros((H, W), dtype=np.float32)
    for lid in PERSON_IDS:
        person_mask[pred == lid] = 1.0

    clothing_mask = np.zeros((H, W), dtype=np.float32)
    for lid in CLOTHING_IDS:
        clothing_mask[pred == lid] = 1.0

    # Dilate person mask to clean up edges
    person_dilated = binary_dilation(person_mask > 0.5, iterations=5).astype(np.float32)
    person_mask = gaussian_filter(person_dilated, sigma=3)

    return person_mask, clothing_mask

def find_holes(alpha_arr, clothing_mask):
    """Find interior holes — areas inside clothing boundary that got erased"""
    clothing_binary = clothing_mask > 0.3
    alpha_binary = alpha_arr > 30

    # Fill the clothing region completely
    filled = binary_fill_holes(clothing_binary)

    # Holes = filled clothing area that has no alpha (person was there)
    holes = filled & ~alpha_binary & clothing_binary == False
    # More robust: holes are inside the bounding box of clothing but have no alpha
    holes = filled & (alpha_arr < 30)
    holes = binary_erosion(holes, iterations=2)
    holes = gaussian_filter(holes.astype(np.float32), sigma=4)
    holes = (holes > 0.3).astype(np.float32)
    holes = gaussian_filter(holes, sigma=3)
    return holes

def inpaint_holes(img_rgb, alpha_u8, holes_mask):
    """Use SD inpainting to fill holes"""
    try:
        pipe = get_inpaint_pipe()
        orig_size = img_rgb.size
        W, H = orig_size

        # Composite clothing on neutral bg for SD context
        neutral = Image.new("RGB", orig_size, (242, 241, 238))
        neutral.paste(img_rgb, mask=Image.fromarray(alpha_u8))

        # Resize to 512 for SD (must be multiple of 8)
        sd_size = (512, 512)
        img_512 = neutral.resize(sd_size, Image.LANCZOS)
        mask_512 = Image.fromarray((holes_mask * 255).astype(np.uint8)).resize(sd_size, Image.LANCZOS)

        result_512 = pipe(
            prompt="white fabric, clothing textile, smooth fabric, seamless, product photography",
            negative_prompt="person, body, skin, face, hands, arms, legs, background, shadow",
            image=img_512,
            mask_image=mask_512,
            num_inference_steps=25,
            guidance_scale=8.0,
        ).images[0]

        result = result_512.resize(orig_size, Image.LANCZOS)
        print("✓ Inpainting done")
        return result
    except Exception as e:
        print(f"Inpainting failed: {e}")
        return img_rgb.convert("RGB")

def normalize_image(img, target_size=(800, 1067)):
    """Resize image to standard size maintaining aspect, pad with white"""
    img_rgb = img.convert("RGB")
    img_rgb.thumbnail(target_size, Image.LANCZOS)
    canvas = Image.new("RGB", target_size, (255, 255, 255))
    x = (target_size[0] - img_rgb.width) // 2
    y = (target_size[1] - img_rgb.height) // 2
    canvas.paste(img_rgb, (x, y))
    return canvas

def process_single(img, do_inpaint=True):
    """Full ghost mannequin pipeline for one image"""
    img_rgb = normalize_image(img)

    # 1. Remove background
    print("Removing background...")
    bg_removed = remove_bg(img_rgb)
    bg_alpha = np.array(bg_removed.split()[3]).astype(np.float32) / 255.0

    # 2. Get person + clothing masks
    print("Segmenting...")
    person_mask, clothing_mask = get_masks(img_rgb)

    # 3. Remove person pixels from alpha
    clean_alpha = np.clip(bg_alpha * (1.0 - person_mask), 0, 1)
    clean_alpha = gaussian_filter(clean_alpha, sigma=1)
    clean_alpha_u8 = (clean_alpha * 255).astype(np.uint8)

    # 4. Find and fill holes
    if do_inpaint:
        print("Finding holes...")
        holes = find_holes(clean_alpha_u8, clothing_mask)
        has_holes = holes.max() > 0.2
        if has_holes:
            print(f"Inpainting holes (max={holes.max():.2f})...")
            inpainted = inpaint_holes(
                bg_removed.convert("RGB"),
                clean_alpha_u8,
                holes
            )
            inpainted_arr = np.array(inpainted.convert("RGB"))
            base_arr = np.array(bg_removed.convert("RGB"))
            h3 = holes[:,:,np.newaxis]
            merged = (base_arr*(1-h3) + inpainted_arr*h3).astype(np.uint8)
            result_rgb = Image.fromarray(merged)
            new_alpha = np.clip(clean_alpha + holes * 0.85, 0, 1)
            new_alpha = gaussian_filter(new_alpha, sigma=1)
        else:
            result_rgb = bg_removed.convert("RGB")
            new_alpha = clean_alpha
    else:
        result_rgb = bg_removed.convert("RGB")
        new_alpha = clean_alpha

    result_rgba = result_rgb.convert("RGBA")
    result_rgba.putalpha(Image.fromarray((new_alpha*255).astype(np.uint8)))
    return result_rgba

def composite_front_back(front_rgba, back_rgba):
    """3D ghost mannequin: front garment + back neck/armhole interior"""
    # Normalize both to same size
    size = (800, 1067)
    if front_rgba.size != size:
        front_rgba = front_rgba.resize(size, Image.LANCZOS)
    if back_rgba.size != size:
        back_rgba = back_rgba.resize(size, Image.LANCZOS)

    W, H = size
    front_arr = np.array(front_rgba.convert("RGBA")).astype(np.float32)
    back_arr = np.array(back_rgba.convert("RGBA")).astype(np.float32)

    front_alpha = front_arr[:,:,3] / 255.0
    back_alpha = back_arr[:,:,3] / 255.0

    # Interior region: top 30% of image (neck/collar area)
    interior_h = int(H * 0.30)

    # Where front has no clothing but back does = show interior
    front_top = front_alpha[:interior_h,:]
    back_top = back_alpha[:interior_h,:]
    show_back = (front_top < 0.2) & (back_top > 0.3)
    show_back_f = gaussian_filter(show_back.astype(np.float32), sigma=4)

    # Build composite
    composite = front_arr.copy()
    for c in range(4):
        composite[:interior_h,:,c] = (
            front_arr[:interior_h,:,c] * (1 - show_back_f) +
            back_arr[:interior_h,:,c] * show_back_f
        )

    # Fix alpha: union of front and interior fill
    comp_alpha = front_alpha.copy()
    comp_alpha[:interior_h,:] = np.maximum(
        front_alpha[:interior_h,:],
        show_back_f * back_alpha[:interior_h,:]
    )
    composite[:,:,3] = (comp_alpha * 255).astype(np.uint8)

    return Image.fromarray(composite.astype(np.uint8))

def place_on_canvas(garment_rgba):
    W, H = 800, 1067
    bg = Image.new("RGBA", (W, H), (242, 241, 238, 255))
    bbox = garment_rgba.getbbox()
    if bbox:
        garment_rgba = garment_rgba.crop(bbox)
    garment_rgba.thumbnail((int(W*0.86), int(H*0.90)), Image.LANCZOS)
    x = (W - garment_rgba.width) // 2
    y = int(H * 0.04)
    shadow = Image.new("RGBA", (W, H), (0,0,0,0))
    s = Image.new("RGBA", garment_rgba.size, (0,0,0,35))
    s.putalpha(garment_rgba.split()[3])
    shadow.paste(s, (x+6, y+10))
    shadow = shadow.filter(ImageFilter.GaussianBlur(14))
    bg = Image.alpha_composite(bg, shadow)
    bg.paste(garment_rgba, (x, y), garment_rgba)
    return bg.convert("RGB")

# ── Routes ─────────────────────────────────────────────────────────────────

@app.route("/process", methods=["POST"])
def process():
    try:
        f = request.files.get("image")
        if not f:
            return jsonify({"error": "No image"}), 400
        mode = request.form.get("mode", "full")
        do_inpaint = request.form.get("inpaint", "true") == "true"
        img = Image.open(f.stream)
        print(f"Single: {img.size}, mode={mode}, inpaint={do_inpaint}")
        if mode == "full":
            garment = process_single(img, do_inpaint=do_inpaint)
        else:
            img_norm = normalize_image(img)
            garment = remove_bg(img_norm)
        result = place_on_canvas(garment)
        buf = io.BytesIO(); result.save(buf, "PNG"); buf.seek(0)
        return send_file(buf, mimetype="image/png")
    except Exception as e:
        print(traceback.format_exc())
        return jsonify({"error": str(e)}), 500

@app.route("/composite", methods=["POST"])
def composite():
    try:
        front_f = request.files.get("front")
        back_f = request.files.get("back")
        if not front_f or not back_f:
            return jsonify({"error": "Need both front and back images"}), 400
        do_inpaint = request.form.get("inpaint", "true") == "true"
        front_img = Image.open(front_f.stream)
        back_img = Image.open(back_f.stream)
        print(f"Composite: front={front_img.size} back={back_img.size}")
        print("Processing front...")
        front_rgba = process_single(front_img, do_inpaint=do_inpaint)
        print("Processing back...")
        back_rgba = process_single(back_img, do_inpaint=do_inpaint)
        print("Compositing...")
        composited = composite_front_back(front_rgba, back_rgba)
        result = place_on_canvas(composited)
        buf = io.BytesIO(); result.save(buf, "PNG"); buf.seek(0)
        print("Done!")
        return send_file(buf, mimetype="image/png")
    except Exception as e:
        print(traceback.format_exc())
        return jsonify({"error": str(e)}), 500

@app.route("/")
def index():
    return open("/app/index.html").read()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)