Update preprocess_onepair.py

preprocess_onepair.py

@@ -1,136 +1,65 @@
 # 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
+import io, os, shutil
 from pathlib import Path
-from typing import Optional, Tuple
-
+from typing import 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)
 
+_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")
+    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))
+def _resize_pad(im: Image.Image, size: Tuple[int,int]=(512,512)) -> Image.Image:
+    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
+    cut = remove(arr)
+    alpha = cut[:,:,3] if cut.shape[2]==4 else np.ones(arr.shape[:2],dtype=np.uint8)*255
+    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
+    res = _mp_seg.process(cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
+    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:
+    root = Path(root); test_root = root/"test"
+    if test_root.exists(): shutil.rmtree(test_root)
+    for sub in ["image","cloth","edge","image-parse","pose","warp_feat"]:
+        (test_root/sub).mkdir(parents=True,exist_ok=True)
 
-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"))
+    person_pil=_resize_pad(_to_pil(person_img)); garment_pil=_resize_pad(_to_pil(garment_img))
+    person_name="user_0001.jpg"; cloth_name="cloth_0001.jpg"
+    person_pil.save(test_root/"image"/person_name,quality=95)
+    garment_pil.save(test_root/"cloth"/cloth_name,quality=95)
 
-    # 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"))
+    edge_pil=_cloth_edge(garment_pil).convert("L").resize((512,512),Image.NEAREST)
+    edge_pil.save(test_root/"edge"/cloth_name.replace(".jpg",".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")
+    parse_pil=_human_mask(person_pil).convert("L")
+    parse_pil.save(test_root/"image-parse"/person_name.replace(".jpg",".png"))
 
-    # pairs.txt
-    with open(test_root / "pairs.txt", "w") as f:
-        f.write(f"{person_name} {cloth_name}\n")
+    dummy=np.zeros((512,512,3),dtype=np.uint8)
+    Image.fromarray(dummy).save(test_root/"pose"/person_name.replace(".jpg","_keypoints.png"))
+    feat=np.zeros((256,256,3),dtype=np.uint8)
+    Image.fromarray(feat).save(test_root/"warp_feat"/f"{person_name[:-4]}_{cloth_name[:-4]}.png")
 
+    with open(test_root/"pairs.txt","w") as f: f.write(f"{person_name} {cloth_name}\n")
     return str(root)