warp_design_on_dress.py

import os
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
from networks import GMM, UnetGenerator, load_checkpoint


def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    # Preprocessing
    im_h, im_w = 256, 192
    tf = transforms.Compose([
        transforms.Resize((im_h, im_w)),
        transforms.ToTensor()
    ])

    dress_img = Image.open(dress_path).convert("RGB")
    design_img = Image.open(design_path).convert("RGB")

    dress_tensor = tf(dress_img).unsqueeze(0).cpu()
    design_tensor = tf(design_img).unsqueeze(0).cpu()
    design_mask = torch.ones_like(design_tensor[:, :1, :, :])  # full white mask

    # Fake agnostic: use the dress image itself
    agnostic = dress_tensor.clone()

    # ----- GMM -----
    gmm = GMM(opt=None)
    load_checkpoint(gmm, gmm_ckpt)
    gmm.cuda().eval()

    with torch.no_grad():
        grid, _ = gmm(agnostic, design_mask)
        warped_design = F.grid_sample(design_tensor, grid, padding_mode='border')
        warped_mask = F.grid_sample(design_mask, grid, padding_mode='zeros')

    # ----- TOM -----
    tom = UnetGenerator(26, 4, 6, ngf=64, norm_layer=torch.nn.InstanceNorm2d)
    load_checkpoint(tom, tom_ckpt)
    tom.cpu().eval()

    with torch.no_grad():
        tom_input = torch.cat([agnostic, warped_design, warped_mask], 1)
        output = tom(tom_input)

        p_rendered, m_composite = torch.split(output, 3, 1)
        p_rendered = torch.tanh(p_rendered)
        m_composite = torch.sigmoid(m_composite)

        tryon = warped_design * m_composite + p_rendered * (1 - m_composite)

        # Save output
        out_img = tryon.squeeze().permute(1, 2, 0).cpu().numpy()
        out_img = (out_img * 255).astype("uint8")
        out_pil = Image.fromarray(out_img)

        output_path = os.path.join(output_dir, "tryon.jpg")
        out_pil.save(output_path)

        return output_path