Create preprocess_onepair.py

preprocess_onepair.py

@@ -0,0 +1,136 @@
+# preprocess_onepair.py
+# Build a one-pair VITON-like dataset from a user photo + garment.
+# Creates: /<root>/test/{image,cloth,edge,image-parse,pose,warp_feat} + pairs.txt
+#
+# NOTES:
+# - Cloth edge is made with rembg (background removal).
+# - Human mask uses MediaPipe Selfie Segmentation (coarse but works).
+# - pose, image-parse, warp_feat are stubbed for now (zeros/placeholders).
+#   You can replace them later with SCHP/OpenPose/PF-AFN for higher quality.
+
+from __future__ import annotations
+import os
+import io
+import shutil
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+from PIL import Image
+import cv2
+
+# Background removal for cloth
+from rembg import remove
+
+# Coarse human mask (person vs background)
+import mediapipe as mp
+_mp_seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)
+
+
+def _to_pil(img) -> Image.Image:
+    if isinstance(img, Image.Image):
+        return img
+    if isinstance(img, (str, os.PathLike)):
+        return Image.open(img).convert("RGB")
+    if isinstance(img, bytes):
+        return Image.open(io.BytesIO(img)).convert("RGB")
+    raise TypeError("Unsupported image type")
+
+
+def _resize_pad(im: Image.Image, size: Tuple[int, int]=(512, 512)) -> Image.Image:
+    """Letterbox to target size, keeping aspect, filled with white."""
+    w, h = im.size
+    tw, th = size
+    scale = min(tw / w, th / h)
+    nw, nh = int(w * scale), int(h * scale)
+    im2 = im.resize((nw, nh), Image.BICUBIC)
+    canvas = Image.new("RGB", (tw, th), (255, 255, 255))
+    canvas.paste(im2, ((tw - nw)//2, (th - nh)//2))
+    return canvas
+
+
+def _cloth_edge(garment_rgb: Image.Image) -> Image.Image:
+    """Remove background → alpha → binary edge (white cloth on black)."""
+    arr = np.array(garment_rgb)
+    cut = remove(arr)  # RGBA
+    if cut.shape[2] == 4:
+        alpha = cut[:, :, 3]
+    else:
+        alpha = np.ones(arr.shape[:2], dtype=np.uint8) * 255
+
+    # Edge is "white where cloth exists"
+    edge = np.zeros_like(alpha, dtype=np.uint8)
+    edge[alpha > 10] = 255
+    return Image.fromarray(edge)
+
+
+def _human_mask(human_rgb: Image.Image) -> Image.Image:
+    """Coarse person mask via MediaPipe (returns white=person, black=bg)."""
+    arr = np.array(human_rgb)
+    rgb = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
+    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+    res = _mp_seg.process(rgb)
+    mask = (res.segmentation_mask > 0.5).astype(np.uint8) * 255
+    return Image.fromarray(mask)
+
+
+def build_temp_dataset(
+    person_img,
+    garment_img,
+    root: Path | str
+) -> str:
+    """
+    Build one-pair dataset under <root>/test and return dataroot as string.
+    """
+    root = Path(root)
+    test_root = root / "test"
+
+    # Clean slate
+    if test_root.exists():
+        shutil.rmtree(test_root)
+    (test_root / "image").mkdir(parents=True, exist_ok=True)
+    (test_root / "cloth").mkdir(parents=True, exist_ok=True)
+    (test_root / "edge").mkdir(parents=True, exist_ok=True)
+    (test_root / "image-parse").mkdir(parents=True, exist_ok=True)
+    (test_root / "pose").mkdir(parents=True, exist_ok=True)
+    (test_root / "warp_feat").mkdir(parents=True, exist_ok=True)
+
+    # Normalize to 512×512 canvases
+    person_pil = _resize_pad(_to_pil(person_img), (512, 512))
+    garment_pil = _resize_pad(_to_pil(garment_img), (512, 512))
+
+    # Save base images
+    person_name = "user_0001.jpg"
+    cloth_name = "cloth_0001.jpg"
+    (test_root / "image" / person_name).parent.mkdir(parents=True, exist_ok=True)
+    (test_root / "cloth" / cloth_name).parent.mkdir(parents=True, exist_ok=True)
+
+    person_pil.save(test_root / "image" / person_name, quality=95)
+    garment_pil.save(test_root / "cloth" / cloth_name, quality=95)
+
+    # Cloth edge (white cloth mask on black)
+    edge_pil = _cloth_edge(garment_pil)
+    edge_pil = edge_pil.convert("L").resize((512, 512), Image.NEAREST)
+    edge_pil.save(test_root / "edge" / cloth_name.replace(".jpg", ".png"))
+
+    # Human parse (stub = coarse person mask)
+    # Most CP datasets expect a PNG with label IDs; we give a soft body mask
+    # where torso is expected to be > 0; downstream code will use it as inpaint mask.
+    parse_pil = _human_mask(person_pil).convert("L")
+    parse_pil.save(test_root / "image-parse" / person_name.replace(".jpg", ".png"))
+
+    # Pose (stub): create empty JSON-like signal as a png (many repos read JSON;
+    # your CPDataset may not require pose if inpaint mask is present)
+    dummy_pose = np.zeros((512, 512, 3), dtype=np.uint8)
+    Image.fromarray(dummy_pose).save(test_root / "pose" / person_name.replace(".jpg", "_keypoints.png"))
+
+    # Warp feat (stub): Some loaders expect precomputed warped feat tensor images.
+    # We drop a zero tile; if your CPDataset ignores it at test time, fine.
+    wfeat = np.zeros((256, 256, 3), dtype=np.uint8)
+    Image.fromarray(wfeat).save(test_root / "warp_feat" / f"{person_name[:-4]}_{cloth_name[:-4]}.png")
+
+    # pairs.txt
+    with open(test_root / "pairs.txt", "w") as f:
+        f.write(f"{person_name} {cloth_name}\n")
+
+    return str(root)