simulate_surfers

#!/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 enum import Enum
from math import pow, cos, sin

import numpy as np
import cv2

import utils


__doc__ = (
    'Click on the image to deploy a new UAV that flies over Galicia and surfs '
    'on the pre-computed field of forces to reach the regions with highest '
    'risk of fire.')


WINDOW_NAME = 'Surfers'
AGENT_BODY = (0, 0, 0, 255)
AGENT_VISION = (0, 0, 0, 50)
FORCE_FACTOR =  0.005
SEPARATION_RANGE = 40
SEPARATION_FACTOR = 0.03
MAX_SPEED = 0.4


class Agent(object):
    def __init__(self, x, y):
        self.position = np.array((x, y))
        self.velocity = np.array((0, 1))


def parse_args():
    """
    Parses map and field of forces.
    """
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('map', type=str, help='map file')
    parser.add_argument('forces', type=str, help='field of forces file')
    return parser.parse_args()


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


def on_mouse_event(event, x, y, flags, agents):
    if event == cv2.EVENT_LBUTTONDOWN:
        agents.append(Agent(x, y))


def update_agents(agents, field_of_forces, cell_size):
    for agent in agents:
        # Force that steers to regions of high risk of fire.
        force_from_field = __get_force_from_field(
            agent, field_of_forces, cell_size)
        # Force that separates from other agents.
        separation_force = __get_separation_force(agent, agents)
        # We weight the forces and add them to the velocity
        total_force = (
            force_from_field * FORCE_FACTOR +
            separation_force * SEPARATION_FACTOR)
        velocity = agent.velocity + total_force
        agent.velocity = __clamp_velocity_to_max_speed(velocity)
        agent.position = agent.position + agent.velocity


def __get_force_from_field(agent, field_of_forces, cell_size):
    x, y = agent.position.astype(int)
    force_from_field = np.zeros(2)
    if x > 0 and y > 0:
        row, col = y // cell_size, x // cell_size
        rows, columns, _ = field_of_forces.shape
        if row < rows and col < columns:
            force_from_field = field_of_forces[row][col]
    return force_from_field


def __get_separation_force(agent, agents):
    separation_force = np.zeros(2)
    for other in agents:
        if other is not agent:
            other_force = agent.position - other.position
            other_magnitude = np.linalg.norm(other_force)
            if other_magnitude < SEPARATION_RANGE:
                other_force /= other_magnitude
                separation_force += other_force;
    separation_magnitude = np.linalg.norm(separation_force)
    if separation_magnitude:
        separation_force /= separation_magnitude
    return separation_force


def __clamp_velocity_to_max_speed(velocity):
    speed = np.linalg.norm(velocity)
    if speed > MAX_SPEED:
        velocity = velocity / speed * MAX_SPEED
    return velocity


def create_agents_image(height, width, agents):
    img = np.zeros((height, width, 4), dtype=np.uint8)
    for agent in agents:
        position = tuple(agent.position.astype(int))
        cv2.circle(img, position, 15, AGENT_VISION, -1)
        cv2.circle(img, position, 4, AGENT_BODY, -1)
        cv2.circle(img, position, 16, AGENT_BODY, 1)
    return img


def read_input(key):
    global MAX_SPEED
    if key == 43:
        MAX_SPEED += 0.1
        print(MAX_SPEED)
    elif key == 45:
        MAX_SPEED -= 0.1
        print(MAX_SPEED)


def main():
    args = parse_args()

    print('Reading {} ...'.format(args.forces))
    field_of_forces = np.load(args.forces)

    print('Reading {} ...'.format(args.map))
    img = utils.read_image_with_alpha(args.map)

    agents = []

    cv2.namedWindow(WINDOW_NAME)
    cv2.setMouseCallback(WINDOW_NAME, on_mouse_event, agents)

    height, width, _ = img.shape
    cell_size = height // field_of_forces.shape[0]
    key = 0
    while key != 27 and utils.window_is_open(WINDOW_NAME):
        read_input(key)
        update_agents(agents, field_of_forces, cell_size)
        agents_img = create_agents_image(height, width, agents)
        key = show_step(utils.blend(img, agents_img), 1)


if __name__ == '__main__':
    main()