from unittest import result
import PIL.Image as Image
import numpy as np
from mmseg.core.evaluation import eval_metrics
from mmseg.apis import inference_segmentor, init_segmentor
import torch
from os import path as osp
import os
import sys

from mmseg.models.segmentors import base

def print_progress(current, total):
    sys.stdout.write('\r')
    sys.stdout.write(f'{current} out of {total} processed.')
    sys.stdout.flush()

def evaluate_custom_model(config_file, checkpoint_file, basedir, person_idx):
    """Evaluate a model by calculating mIoU with a given dataset
    Args:
        basedir (str): Images and annotations are in [basedir]/images/test, [basedir]/annotations/test
        person_idx (int): 1 for our person_mix, 11 for Cityscapes
    """

    # config_file = 'configs/bisenetv2/bisenetv2_fcn_4x4_1024x1024_160k_cityscapes.py'
    # checkpoint_file = 'pretrained/bisenetv2_fcn_4x4_1024x1024_160k_cityscapes_20210902_015551-bcf10f09.pth'
    # config_file = 'configs/stdc/stdc1_512x512_80k_custom.py'
    # checkpoint_file = 'work_dirs/stdc1_512x512_80k_custom/models/v7_best.pth'

    # build the model from a config file and a checkpoint file
    model = init_segmentor(config_file, checkpoint_file)

    jpg_files_base = [filename for filename in os.listdir(
        basedir + 'images/test') if filename.lower().endswith('.jpg')]
    jpg_files = [basedir + 'images/test/' + filename for filename in jpg_files_base]
    png_files = [basedir + 'annotations/test/' + osp.splitext(filename)[0] + '.png' for filename in jpg_files_base]

    results = []
    counter = 0
    for img_file in jpg_files:
        result = inference_segmentor(model, img_file)[0] # Only one image

        # Convert person pixels to 0 while the rest to 1
        if person_idx != 1:
            result[result != person_idx] = 0
            result[result == person_idx] = 1

        results.append(result)

        # Print progress
        counter += 1
        print_progress(counter, len(jpg_files))

    model.show_result(jpg_files[0], results, show=True)

    eval_results = eval_metrics(results, png_files, 2, ignore_index=255)
    print('Person mIoU: ' + str(eval_results['IoU'][1]))


if __name__ == '__main__':
    basedir = osp.expanduser('./BinaryAnnotationsMIXv7/')
    cfg_file = sys.argv[1]
    ckp_file = sys.argv[2]

    person_idx = 1
    if 'cityscapes' in cfg_file:
        person_idx = 11
    elif 'ade20' in cfg_file:
        person_idx = 12

    evaluate_custom_model(
        cfg_file,
        ckp_file,
        basedir,
        person_idx)
