# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Calculate tile classes per H5 file
#
# Once the Ilastik classifier has classified images, it produces per-pixel probabilities as an H5 file. This workbook loads those per-pixel probabilities, and then samples each tile to obtain the class for all the pixels in the tile, so that the tile can be compared with the ground truth data as a whole

# ## Step 1: Determine how many pixels to sample per tile for a statistically significant finding
#
# Use the Cochran formula to determine a statistically significant sample of pixels. Then trim that down to the total population of a texture tile (900 pixels)

def sample_size(tile_size):
    # Calculate statistically significant test sample size

    # Z score for a 95% confidence (from Z table)
    Z_score = 1.96
    # Calculate with a 5% margin of error
    margin_of_error = 0.05


    # Some selected expected population percentages
    population_with_attribute_water = 0.09
    population_with_attribute_foliage = 0.876
    population_with_attribute_road = 0.027
    population_with_attribute_building = 0.006
    
    # Number of samples is the size of a square tile
    N = tile_size ** 2

    required_samples = 0

    # Calc for all population percentages, output all values. We will pick the largest one
    for p in [population_with_attribute_water, 
              population_with_attribute_foliage, population_with_attribute_road, 
              population_with_attribute_building]:
        q = 1 - p
        # Calc required samples for an unlimited population
        n_0 = ((Z_score ** 2) * p * q) / (margin_of_error ** 2)
        # Now reduce this unbounded required sample count down by the known population 
        # per tile (900 pixels)
        n = n_0 / (1 + ((n_0 - 1) / N))
        if (n > required_samples):
            required_samples = round(n) + (1 if round(n) != n else 0)

    return required_samples

# ## Step 2: Calculate the class per tile in each probability file, save as csv
#
# Each image will have a statistically significant sample of pixels chosen per tile (uniformly distributed), and the probabilities of those pixels will be aggregated to provide the class of the tile.
#
# A CSV file will be saved per image with the class of all the tiles in the image

# +
import numpy as np
import matplotlib.pyplot as plt
import h5py
import pandas as pd
import itertools
import os
import glob

def calc_class(foliage, water, building, road):
    if foliage > 0.5:
        return 'foliage'
    if water > 0.5:
        return 'water'
    if building > 0.5:
        return 'building'
    if road > 0.5:
        return 'road'
    return 'unknown'

def calc_tile_classes(img_data_arr, tile_size):    
    # The image data is in a 3 dimensional array. Dimension 1 is the Y pixel value, 2 is the X pixel value 
    # and Z contains an array of the 4 classes probabilities
    names = ['y', 'x', 'z']
    # Create an index for the dataframe
    index = pd.MultiIndex.from_product([range(s) for s in img_data_arr.shape], names=names)
    # create the dataframe itself
    image_df = pd.DataFrame({'A': img_data_arr.flatten()}, index=index)['A']
    # Reformat into a 4 column frame with the 2 column index by unpacking the array
    image_df = image_df.unstack(level='z').swaplevel().sort_index()
    # Set the column names for the probabilities
    image_df.columns = ['A', 'B', 'C', 'D']

    y_size, x_size, z_size = img_data_arr.shape
    required_samples = sample_size(tile_size)
    # Now build a matrix of sample pixels across the image, the array will contain the 30 pixel tile number
    # and a uniformly distributed random selection of pixels from that tile
    sample_pixel_arr = []
    all_items = itertools.product(np.asarray(range(int(x_size / tile_size))), np.asarray(range(int(y_size/tile_size))))
    for x, y in all_items:
        x_s = np.asarray([int(round(x[0])) for x in np.random.uniform((x * tile_size), (x * tile_size) + 
                                                                      tile_size, (required_samples, 1))])
        y_s = np.asarray([int(round(x[0])) for x in np.random.uniform((y * tile_size), (y * tile_size) + 
                                                                      tile_size, (required_samples, 1))])
        sample_pixel_arr.append(list(zip(x_s, y_s, itertools.repeat(x), itertools.repeat(y))))

    # Convert the sample array to a dataframe for joining
    sample_matrix = np.asarray(sample_pixel_arr)
    sample_arr = np.reshape(sample_matrix, (sample_matrix.shape[0] * sample_matrix.shape[1], 
                                            sample_matrix.shape[2]))
    sample_df = pd.DataFrame(sample_arr)
    sample_df.columns = ['x', 'y', 'tile_x', 'tile_y']
    sample_df.set_index(['x', 'y'], inplace=True)
    
    # Join the sample pixels with the original probabilities frame to get the probabilities for each sample pixel
    sample_pixels = pd.merge(left=sample_df, right=image_df, left_on=['x', 'y'], right_on=['x', 'y'])

    # Sum all probabilities to give a total probability value for each category
    aggregated_samples = sample_pixels.groupby(['tile_x', 'tile_y'], as_index=False).agg(
        {"A": "sum", "B": "sum", "C": "sum", "D": "sum"})
    
    aggregated_samples = aggregated_samples.fillna(0)
    
    sample_counts = [0] * len(aggregated_samples)        
    sample_counts = sample_counts + aggregated_samples['A']
    sample_counts = sample_counts + aggregated_samples['B']
    sample_counts = sample_counts + aggregated_samples['C']
    sample_counts = sample_counts + aggregated_samples['D']   
        
    # Divide total probability by number of samples to give a probability percentage per tile
    aggregated_samples['A'] = aggregated_samples['A'] / sample_counts
    aggregated_samples['B'] = aggregated_samples['B'] / sample_counts
    aggregated_samples['C'] = aggregated_samples['C'] / sample_counts
    aggregated_samples['D'] = aggregated_samples['D'] / sample_counts    
    
    # Now calculate the class by inspecting the probability percentage. If any of the categories is 
    # above 50% that is considered to be the predicted  category of that tile
    aggregated_samples.loc[:,'tile_class'] = pd.Series('unsure', index=aggregated_samples.index)
    aggregated_samples.loc[aggregated_samples['A'] > 0.5, 'tile_class'] = pd.Series('foliage', index=aggregated_samples.index)
    aggregated_samples.loc[aggregated_samples['B'] > 0.5, 'tile_class'] = pd.Series('water', index=aggregated_samples.index)
    aggregated_samples.loc[aggregated_samples['C'] > 0.5, 'tile_class'] = pd.Series('building', index=aggregated_samples.index)
    aggregated_samples.loc[aggregated_samples['D'] > 0.5, 'tile_class'] = pd.Series('road', index=aggregated_samples.index)
    #aggregated_samples['tile_class'] = df.apply(lambda x: calc_class(x['A'], x['B'], x['C'], x['D']), axis=1)

    return aggregated_samples


# +
import re 
from PIL import Image

probability_filenames = glob.glob('../../ilastik/TestingData/*/*/*.h5', recursive=True)
for file in probability_filenames:
    head_tail = os.path.split(file)
    folder = os.path.basename(head_tail[0])
    filename = head_tail[1]
    pre, ext = os.path.splitext(filename)
    
    image_file_portion = re.search("DJI\_[0-9]*", pre).group()
    search_folder = '../../Texture_Repo/Donegal_Rural_Terrain_Textures/Test_Images/*/' + image_file_portion + '.jpg'
    match_files = glob.glob(search_folder)
    if len(match_files) == 0:
        print('ERROR: Source file not found for', pre)
    else:
        orig_image = Image.open(match_files[0])
        original_width = orig_image.width
        f = h5py.File(file, 'r')
        img_data_arr = np.asarray(f['exported_data'])                
        classified_width = img_data_arr.shape[1]
        tile_size = (classified_width / original_width) * 30
        print('Processing', file, 'orig width', original_width, 'classified width', classified_width, 
              'tile size', tile_size)    
        tile_data = calc_tile_classes(img_data_arr, tile_size)
        tile_data.to_csv('../../TestPredictions/Ilastik/' + folder + '_' + pre + '.csv', index=False)    

StarcoderdataJupyterScriptsDedupFiltered