import sys
import os
import cv2
import torch
import warnings
import numpy as np
import albumentations as albs
import segmentation_models_pytorch as smp
from albumentations.pytorch.transforms import ToTensorV2
from torch import nn


MODEL_PATH = "model_final.pth"
IMG_SIZE = 256
FRAME_SIZE = 1024


def predict(image, model, device):
    image_shape = image.shape
    alb_transforms = albs.Compose([albs.Resize(IMG_SIZE, IMG_SIZE), albs.Normalize(), ToTensorV2()])
    now = alb_transforms(image=image)['image'].to(device).view(1, 3, IMG_SIZE, IMG_SIZE)
    piction = model(now)
    piction = nn.Sigmoid()(piction)
    piction = (piction > 0.45)
    piction = piction[0].detach().cpu().numpy().astype(float)[0]
    piction = albs.Resize(image_shape[0], image_shape[1])(image=piction)['image']

    return piction.astype(bool)


def get_model():
    return smp.Unet(encoder_name="efficientnet-b6", encoder_weights=None, in_channels=3, classes=1)


def get_shape(shape):
    return (shape[0] + FRAME_SIZE - 1) // FRAME_SIZE * FRAME_SIZE, (shape[1] + FRAME_SIZE - 1) // FRAME_SIZE * FRAME_SIZE, shape[2]


def main():
    device = torch.device("cuda")
    tips, output_path = sys.argv[1:]
    model = get_model().to(device)
    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    model = model.eval()
    results = {}

    for i in os.listdir(tips):
        img_path = os.path.join(tips, i)
        im = cv2.imread(img_path)
        image = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
        mask = predict(image, model, device)
        results[i] = mask
    np.savez_compressed(os.path.join(output_path), **results)


if __name__ == "__main__":
    main()