generate_field_of_forces

#!/usr/bin/env python

#
#    Fire Detection Unmanned Aerial System
#    Copyright (C) 2019 Carlos Perez-Lopez
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

import argparse
from math import pow, cos, sin

import numpy as np
import cv2

import utils


__doc__ = 'Generates a field of forces based on Galicia\'s IRDI map.'


WINDOW_NAME = 'Generate Field of Forces'
MAX_SMOOTH_ITERATIONS = 1


def parse_args():
    """
    Parses map and cell size.
    """
    parser = argparse.ArgumentParser( description=__doc__)
    parser.add_argument('map', type=str, help='map file')
    parser.add_argument('cell_size', type=int, help='cell size')
    return parser.parse_args()


def show_step(image, delay):
    utils.show_step(WINDOW_NAME, image, delay)


def calculate_forces_on_borders(cells):
    rows, columns = cells.shape
    forces = np.zeros((rows, columns, 2), dtype=float)
    row_steps = (-1, -1, -1,  0, 0,  1, 1, 1);
    col_steps = (-1,  0,  1, -1, 1, -1, 0, 1);
    for i in range(rows):
        for j in range(columns):
            risk = cells[i][j]
            neighbours = {}
            for ii, jj in zip(row_steps, col_steps):
                row = (i + ii) % rows
                column = (j + jj) % columns
                neighbour_risk = cells[row][column]
                if neighbour_risk in neighbours:
                    neighbours[neighbour_risk].append(
                        np.array((jj, ii)))
                else:
                    neighbours[neighbour_risk] = [np.array((jj, ii))]
            forces[i][j] = __get_gradient_force(risk, neighbours)
    return forces


def __get_gradient_force(risk, neighbours):
    """
    Calculates the vector that points to the maximum risk.
    """
    if len(neighbours) == 1 and risk in neighbours:
        # There is no gradient, all cells are in the same region.
        return np.zeros(2)

    highest_risk = max(neighbours.keys())
    if highest_risk < risk:
        # Only touching lower risk cells.
        return np.zeros(2)

    gradient_force = np.zeros(2)
    for force in neighbours[highest_risk]:
        gradient_force += force
    magnitude = np.linalg.norm(gradient_force)
    if magnitude:
        return gradient_force / magnitude
    return np.zeros(2)


def create_forces_image(forces, cell_size):
    rows, columns, _ = forces.shape
    img = np.zeros((rows * cell_size, columns * cell_size, 4), dtype=np.uint8)
    for i in range(rows):
        for j in range(columns):
            force = forces[i][j]
            if np.count_nonzero(force):
                force = force * 1/2 * cell_size
                center_x = j * cell_size + cell_size / 2
                center_y = i * cell_size + cell_size / 2
                center = np.array((center_x, center_y), dtype=int)
                point1 = center - 1/2 * force
                point2 = center + 1/2 * force
                point1 = tuple(point1.astype(int))
                point2 = tuple(point2.astype(int))
                cv2.arrowedLine(
                    img, point1, point2, (0, 0, 0, 255), 1,
                    line_type=cv2.LINE_8, tipLength=.25)
    return img


def grow_field_of_forces(forces, min_influence=3):
    grown_forces = np.zeros(forces.shape)
    rows, columns, _ = forces.shape
    row_steps = (-1, -1, -1,  0, 0,  1, 1, 1);
    col_steps = (-1,  0,  1, -1, 1, -1, 0, 1);
    for i in range(rows):
        for j in range(columns):
            force = forces[i][j]
            if np.count_nonzero(force):
                grown_forces[i][j] = force
            else:
                neighbours = []
                for ii, jj in zip(row_steps, col_steps):
                    row = (i + ii) % rows
                    column = (j + jj) % columns
                    neighbour_force = forces[row][column]
                    if np.count_nonzero(neighbour_force):
                        neighbours.append(neighbour_force)
                if len(neighbours) >= min_influence:
                    grown_forces[i][j] = __get_mean_force(neighbours)
    return grown_forces


def __get_mean_force(forces):
    mean_force = np.zeros(2)
    for force in forces:
        mean_force += force
    if np.count_nonzero(mean_force):
        return mean_force / np.linalg.norm(mean_force)
    return np.zeros(2)


def all_forces_are_non_zero(forces):
    rows, columns, _ = forces.shape
    for i in range(rows):
        for j in range(columns):
            if not np.count_nonzero(forces[i][j]):
                return False
    return True


def main():
    args = parse_args()

    print('Reading {} ...'.format(args.map))
    img = utils.read_image_with_alpha(args.map)
    utils.show_step(WINDOW_NAME, img, 0)

    print('Adding grid ...')
    height, width, _ = img.shape
    grid = utils.create_grid(height, width, args.cell_size)
    show_step(utils.blend(img, grid), 0)

    print('Averaging cells ...')
    avg_cells = utils.average_cells(img, args.cell_size)
    img = utils.create_rasterized_image(avg_cells, args.cell_size)
    height, width, _ = img.shape
    grid = utils.create_grid(height, width, args.cell_size)
    show_step(utils.blend(img, grid), 0)

    print('Classifying cells ...')
    risk_cells = utils.classify_cells(avg_cells)
    color_cells = utils.risk_cells_to_color_cells(risk_cells)
    img = utils.create_rasterized_image(color_cells, args.cell_size)
    show_step(utils.blend(img, grid), 0)

    print('Smoothing cells ...')
    for i in range(MAX_SMOOTH_ITERATIONS):
        print(' - iteration ', i)
        smoothed_risk_cells = utils.smooth(risk_cells, 1)
        color_cells = utils.risk_cells_to_color_cells(smoothed_risk_cells)
        img = utils.create_rasterized_image(color_cells, args.cell_size)
        show_step(utils.blend(img, grid), 1)
        if np.array_equal(smoothed_risk_cells, risk_cells):
            break
        risk_cells = smoothed_risk_cells
    show_step(utils.blend(img, grid), 0)

    print('Calculating forces on borders ...')
    forces = calculate_forces_on_borders(risk_cells)
    forces_img = create_forces_image(forces, args.cell_size)
    show_step(utils.blend(img, utils.blend(grid, forces_img)), 0)

    print('Growing field of forces ...')
    for i in range(100):
        print(' - iteration ', i)
        forces = grow_field_of_forces(forces)
        forces_img = create_forces_image(forces, args.cell_size)
        show_step(utils.blend(img, utils.blend(grid, forces_img)), 1)
        if all_forces_are_non_zero(forces):
            break
    show_step(utils.blend(img, utils.blend(grid, forces_img)), 0)

    print('Saving field of forces ...')
    filename = 'field_of_forces.npy'
    np.save(filename, forces)

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()