floatingforests

import numpy as np
import skimage.color as skcolor
import skimage.io as skio
import scipy.stats as stats
import os


def hsv_thresholding(filename):
    im = skio.imread(filename)
    im = skcolor.rgb2hsv(im)

    im[:, :, 0] = stats.threshold(im[:, :, 0], threshmin=.5, threshmax=.7, newval=0)     # hue
    im[:, :, 1] = stats.threshold(im[:, :, 1], threshmin=.05, threshmax=.45, newval=0)   # sat
    im[:, :, 2] = stats.threshold(im[:, :, 2], threshmin=.005, threshmax=.55, newval=0)  # value
    return im


def coast_predicate(filename):

    def count_thresholded_pixels(array):
        flat = np.reshape(np.ravel(array), (array.size//3, 3))
        output = flat[(flat[:, 0] != 0) & (flat[:, 1] != 0) & (flat[:, 2] != 0)]
        return output.shape[0]

    im = hsv_thresholding(filename)

    total_pixels = np.size(im)//3
    water_percent = count_thresholded_pixels(im)/total_pixels
    if verbose: print(round(water_percent, 3))
    if water_percent > 0.05:
        return True

    water_edge = False
    edges = (count_thresholded_pixels(im[0, :, :]), count_thresholded_pixels(im[-1, :, :]),
             count_thresholded_pixels(im[:, 0, :]), count_thresholded_pixels(im[:, -1, :]))
    pixel_threshold = 20

    if any(edge > pixel_threshold for edge in edges):
        if verbose: print('has water edge')
        water_edge = True

    if water_edge and water_percent > 0.01:
        return True


if __name__ == "__main__":
    verbose = False
    datapath = 'data'
    if not os.path.exists('pass'):
        os.mkdir('pass')
    if not os.path.exists('fail'):
        os.mkdir('fail')

    filelist = [os.path.join(datapath, f) for f in os.listdir(datapath) if '.jpg' in f]
    for filename in filelist:
        if verbose: print(filename)

        if coast_predicate(filename):
            if verbose: print('good image')
            os.rename(filename, filename.replace(datapath, 'pass'))
        else:
            os.rename(filename, filename.replace(datapath, 'fail'))
        if verbose: print('----')