|
|
|
|
|
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) |
