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)