File size: 1,188 Bytes
8c9048a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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
# from pylab import imshow
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) |