# -- Code Cell --
import pandas as pd
from PIL import Image
import os
import numpy as np

def image_to_rle(image_path):
    image = Image.open(image_path)
    side, side = image.size
    arr = np.array(image)
    red = arr[:, :, 0].reshape(side * side, order = 'F')
    green = arr[:, :, 1].reshape(side * side, order = 'F')
    blue = arr[:, :, 2].reshape(side * side, order = 'F')
    
    start = 0
    length = 0
    res = []
    for i in range(side * side):
        if red[i] == 255 and green[i] == 255 and blue[i] == 255:
            if length == 0:
                start = i + 1
            length = length + 1
        else:
            if length > 0:
                res.append(start)
                res.append(length)
            length = 0
    if length > 0:
        res.append(start)
        res.append(length)
        
    return res

image_names = []
ground_truth = []

path_prefix = 'HotSpotsDataset/groundTruth_subtask'
for subtask_id in range(1, 5):
    dir_path = path_prefix + str(subtask_id)
    for image_name in os.listdir(dir_path):
        image_names.append(image_name)
        ground_truth.append(image_to_rle(dir_path + '/' + image_name))

# -- Code Cell --
df = pd.DataFrame({'subtaskID' : map(lambda x : 1 + int(x / 250), list(range(1000))),
                   'datapointID' : image_names, 
                   'answer' : ground_truth})
df.to_csv('ground_truth.csv', index = False)