| |
| from __future__ import annotations |
| import io, os, shutil |
| from pathlib import Path |
| from typing import Tuple |
| import numpy as np |
| from PIL import Image |
| import cv2 |
| from rembg import remove |
| 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: |
| 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: |
| arr = np.array(garment_rgb) |
| 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: |
| arr = np.array(human_rgb) |
| 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) |
|
|
| 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) |
|
|
| edge_pil=_cloth_edge(garment_pil).convert("L").resize((512,512),Image.NEAREST) |
| edge_pil.save(test_root/"edge"/cloth_name.replace(".jpg",".png")) |
|
|
| parse_pil=_human_mask(person_pil).convert("L") |
| parse_pil.save(test_root/"image-parse"/person_name.replace(".jpg",".png")) |
|
|
| 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) |
