|
|
from people_segmentation.pre_trained_models import create_model
|
|
|
from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad
|
|
|
from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image
|
|
|
|
|
|
import numpy as np
|
|
|
import cv2
|
|
|
import torch
|
|
|
import albumentations as albu
|
|
|
import os
|
|
|
from tqdm import tqdm
|
|
|
import webdataset as wds
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
|
|
|
dataset = wds.WebDataset("../../../datasets/cc3m/cc3m_attempt12/00{000..332}.tar").decode("rgb8").to_tuple("jpg;png", "json")
|
|
|
|
|
|
root = "/home/kis/Downloads/syn/"
|
|
|
|
|
|
model = create_model("Unet_2020-07-20").cuda()
|
|
|
model.eval()
|
|
|
|
|
|
|
|
|
for fname in os.listdir(root):
|
|
|
image = sample[0]
|
|
|
name = sample[1]["key"]
|
|
|
|
|
|
transform = albu.Compose([albu.Normalize(p=1)], p=1)
|
|
|
padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
|
|
|
x = transform(image=padded_image)["image"]
|
|
|
x = torch.unsqueeze(tensor_from_rgb_image(x), 0).cuda()
|
|
|
|
|
|
with torch.no_grad():
|
|
|
prediction = model(x)[0][0]
|
|
|
|
|
|
mask = (prediction > 0).cpu().numpy().astype(np.uint8)
|
|
|
|
|
|
mask = unpad(mask, pads) * 255
|
|
|
|
|
|
|
|
|
cv2.imwrite(f"output/person_masks/{name}.png", mask) |
