# Check batches

import numpy as np

import pandas as pd

from glob import glob

from tqdm import tqdm

tqdm.pandas()

from PIL import Image

import random

import sys

import os

import cv2

import pdb

import shutil

def batch_check(dataset, num_batches):

    data = pd.read_csv(dataset)

    batch_list = random.sample(range(data['batch'].min(), data['batch'].max()+1), num_batches)

    # Make masks folder path if it doesn't already exist

    if not os.path.exists("batch_check"):

        os.mkdir("batch_check")

    for b in batch_list:

        subset = data[data['batch']==b]

        subset['new_name'] = data['case'] + "_" + data['day'] + "_slice" + data['slice_id'].astype(str)

        subset = subset[['image_path', 'new_name', 'mask_path']]

        path = os.path.join("batch_check", "batch" + str(b))

        if not os.path.exists(path):

            os.makedirs(path)

        else:

            shutil.rmtree(path)           # Removes all the subdirectories!

            os.makedirs(path)

        empty_count = 0 

        for _, row in subset.iterrows():

            image, name, mask_path = row

            # save images in new folder

            shutil.copy(image, os.path.join(path, name + ".png"))

            # check if mask is all black

            mask = np.load(mask_path)

            if mask.sum() == 0:

                empty_count += 1

        empty_pct = 100*empty_count / len(subset)

        print("For batch {}, {}% of masks are empty".format(b, empty_pct))

if __name__ == '__main__':

    # usage: python batch_check.py dataset_name.csv num_batches

    dataset = sys.argv[1]

    num_batches = int(sys.argv[2])

    batch_check(dataset, num_batches)