|
import glob |
|
import os |
|
from pathlib import Path |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn.functional as F |
|
from PIL import Image |
|
from torchvision import transforms |
|
from tqdm import tqdm |
|
import joblib |
|
from huggingface_hub import hf_hub_download |
|
|
|
from .load_u2net import load_cloth_segm_model |
|
from .utils import NormalizeImage, naive_cutout, resize_by_bigger_index, image_resize |
|
|
|
|
|
def segment_garment(inputs_dir, outputs_dir, cls="all"): |
|
os.makedirs(outputs_dir, exist_ok=True) |
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
transform_fn = transforms.Compose( |
|
[transforms.ToTensor(), |
|
NormalizeImage(0.5, 0.5)] |
|
) |
|
|
|
|
|
file_path = hf_hub_download(repo_id="tryonlabs/u2net-cloth-segmentation", filename="u2net_cloth_segm.pth") |
|
|
|
print("model loaded from huggingface:", file_path) |
|
|
|
net = load_cloth_segm_model(device, file_path, in_ch=3, out_ch=4) |
|
|
|
images_list = sorted(os.listdir(inputs_dir)) |
|
pbar = tqdm(total=len(images_list)) |
|
|
|
for image_name in images_list: |
|
img = Image.open(os.path.join(inputs_dir, image_name)).convert('RGB') |
|
img_size = img.size |
|
img = img.resize((768, 768), Image.BICUBIC) |
|
image_tensor = transform_fn(img) |
|
image_tensor = torch.unsqueeze(image_tensor, 0) |
|
|
|
with torch.no_grad(): |
|
output_tensor = net(image_tensor.to(device)) |
|
output_tensor = F.log_softmax(output_tensor[0], dim=1) |
|
output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1] |
|
output_tensor = torch.squeeze(output_tensor, dim=0) |
|
output_arr = output_tensor.cpu().numpy() |
|
|
|
if cls == "all": |
|
classes_to_save = [] |
|
|
|
|
|
for cls in range(1, 4): |
|
if np.any(output_arr == cls): |
|
classes_to_save.append(cls) |
|
elif cls == "upper": |
|
classes_to_save = [1] |
|
elif cls == "lower": |
|
classes_to_save = [2] |
|
elif cls == "dress": |
|
classes_to_save = [3] |
|
else: |
|
raise ValueError(f"Unknown cls: {cls}") |
|
|
|
for cls1 in classes_to_save: |
|
alpha_mask = (output_arr == cls1).astype(np.uint8) * 255 |
|
alpha_mask = alpha_mask[0] |
|
alpha_mask_img = Image.fromarray(alpha_mask, mode='L') |
|
alpha_mask_img = alpha_mask_img.resize(img_size, Image.BICUBIC) |
|
alpha_mask_img.save(os.path.join(outputs_dir, f'{image_name.split(".")[0]}_{cls1}.jpg')) |
|
|
|
pbar.update(1) |
|
|
|
pbar.close() |
|
|
|
|
|
def extract_garment(inputs_dir, outputs_dir, cls="all", resize_to_width=None): |
|
os.makedirs(outputs_dir, exist_ok=True) |
|
cloth_mask_dir = os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask") |
|
os.makedirs(cloth_mask_dir, exist_ok=True) |
|
|
|
segment_garment(inputs_dir, os.path.join(Path(outputs_dir).parent.absolute(), "cloth-mask"), cls=cls) |
|
|
|
images_path = sorted(glob.glob(os.path.join(inputs_dir, "*"))) |
|
masks_path = sorted(glob.glob(os.path.join(cloth_mask_dir, "*"))) |
|
img = Image.open(images_path[0]) |
|
|
|
for mask_path in masks_path: |
|
mask = Image.open(mask_path) |
|
|
|
cutout = np.array(naive_cutout(img, mask)) |
|
cutout = resize_by_bigger_index(cutout) |
|
|
|
canvas = np.ones((1024, 768, 3), np.uint8) * 255 |
|
y1, y2 = (canvas.shape[0] - cutout.shape[0]) // 2, (canvas.shape[0] + cutout.shape[0]) // 2 |
|
x1, x2 = (canvas.shape[1] - cutout.shape[1]) // 2, (canvas.shape[1] + cutout.shape[1]) // 2 |
|
|
|
alpha_s = cutout[:, :, 3] / 255.0 |
|
alpha_l = 1.0 - alpha_s |
|
|
|
for c in range(0, 3): |
|
canvas[y1:y2, x1:x2, c] = (alpha_s * cutout[:, :, c] + |
|
alpha_l * canvas[y1:y2, x1:x2, c]) |
|
|
|
|
|
if resize_to_width: |
|
canvas = image_resize(canvas, width=resize_to_width) |
|
|
|
canvas = Image.fromarray(canvas) |
|
|
|
canvas.save(os.path.join(outputs_dir, f"{os.path.basename(mask_path).split('.')[0]}.jpg"), format='JPEG') |
|
|
