text:
"""Preprocesses AIS data.

To obtain the needed AIS data, see (and run) ``get_raw.sh``.

Uses information specified in the ``config_file`` to chew through available AIS csv data to generate an output csv file
with the discretized states, inferred actions, and records metadata for further processing in the ``meta_file``.
"""
import yaml
import os
import math
import numpy as np
import pandas as pd

def main():
    """Driver code to run the big steps of pre-processing the data.

    First, all the script parameters are loaded by reading the ``.yaml`` ``config_file`` and this config is unpacked.
    All the csv files within a specified directory are found and returned as ``csv_files``, along with each file's
    year, month, and zone in ``all_files_meta``.

    Then, these csv files are read and organized into a large set of trajectories ordered by id (mmsi). Finally, these
    trajectories are discretized before being written into an output csv containing only rows of id-state-action-state
    transitions.

    Another yaml file is written to ``meta_file`` to specify the final grid parameters, output directories, and the
    year, month, and zone of all the files read in.
    """
    # file containing important options, directories, parameters, etc.
    config_file = "config.yml"

    # file to write final grid_params and the csv files' respective years, months, and zones
    meta_file = "meta_data.yml"

    # gets the config dictionary and unpacks it
    config = get_config(config_file)
    options = config["options"]
    directories = config["directories"]
    meta_params = config["meta_params"]
    grid_params = config["grid_params"]

    # gets the csv files available and their metadata
    csv_files, all_files_meta = collect_csv_files(
        options, directories, meta_params
    )

    # reads the collected csv files and assembles trajectories
    trajectories, grid_params = read_data(csv_files, options, grid_params)

    # processes (fits to grid) trajectories and writes generates sequences to output file
    write_data(trajectories, options, directories, grid_params)

    # writes file metadata, paths, and grid parameters to ``meta_file``
    directories_out = {
        "in_dir_path": directories["out_dir_path"],
        "in_dir_data": directories["out_dir_file"],
    }
    out_dict = {
        "all_files_meta": all_files_meta,
        "options": options,
        "directories": directories_out,
        "grid_params": grid_params,
    }
    with open(meta_file, "w") as outfile:
        yaml.dump(out_dict, outfile, default_flow_style=False)

def get_config(config_file):
    """Helper function to get dictionary of script parameters.

    Mostly boilerplate code to read in ``config_file`` as a ``.yaml`` file that specifies important script parameters.
    Upon success, this config dictionary is returned to main to be unpacked.

    Args:
        config_file (str): The name of the ``.yaml`` file containing the script configuration parameters, located in the
            directory the script is run.

    Returns:
        dict: The script configuration parameters.
    """
    with open(config_file, "r") as stream:
        try:
            return yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            print(exc)

def collect_csv_files(options, directories, meta_params):
    """Traverses the directory containing the decompressed AIS data to get the csv names for further processing.

    Uses the ``os`` library to find all csv files within the directory defined by ``directories``, populating the
    ``csv_files`` list for later reading of all valid csv files found and logging file metadata in ``all_files_meta``.

    Args:
        options (dict): The script options specified in the ``config_file``.
        directories (dict): The input and output paths and files specified in the ``config_file``.
        meta_params (dict): The time and zone boundaries specified in the ``config_file``.

    Returns:
        tuple: A list with paths to all valid csv files found and a dictionary with the year, month, and zone
        corresponding to each csv file's origin.
    """
    # initialize a list that will be filled with all csv file names from root
    csv_files = []
    all_files_meta = {}
    for root, dirs, files in os.walk(
        directories["in_dir_path"] + directories["in_dir_data"]
    ):
        for file in files:
            if file.endswith(".csv"):
                year, month, zone = get_meta_data(
                    file
                )  # finds the data year, month, and zone based on the file name

                # only considers valid years and months if time is bounded and valid zones if zones are bounded
                if (
                    not options["bound_time"]
                    or (
                        meta_params["min_year"]
                        <= year
                        <= meta_params["max_year"]
                    )
                ) and (
                    not options["bound_zone"]
                    or (
                        meta_params["min_zone"]
                        <= zone
                        <= meta_params["max_zone"]
                    )
                ):
                    if (
                        not options["bound_time"]
                        or (
                            not year == meta_params["min_year"]
                            or month >= meta_params["min_month"]
                        )
                    ) and (
                        not options["bound_time"]
                        or (
                            not year == meta_params["max_year"]
                            or month <= meta_params["max_month"]
                        )
                    ):
                        # csv_files will contain file locations relative to current directory
                        csv_files.append(os.path.join(root, file))

                        # create dictionary to describe file characteristics
                        file_meta = {
                            "year": year,
                            "month": month,
                            "zone": zone,
                        }
                        all_files_meta[file] = file_meta

    return csv_files, all_files_meta

def read_data(csv_files, options, grid_params):
    """Iterate through each csv file to segregate each trajectory by its mmsi id.

    Reads each csv in ``csv_files`` to obtain coordinates and timestamp series associated with each mmsi id encountered.

    Optionally, the boundaries of the grid later specified can be inferred by calculating the minimum and maximum
    longitudes and latitudes by setting ``options['bound_lon']`` and ``options['bound_lat']`` to ``False``, respectively
    in the ``config_file``.

    It can also be specified to only read the first ``options['max_rows']`` of each csv file by setting
    ``options['limit_rows']`` to True in the ``config_file``.

    Args:
        csv_files (list): Paths to all valid csv files found.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        tuple: A pandas DataFrame of all data entries with the format ``['MMSI', 'LON', 'LAT', 'TIME']`` and a
        dictionary that specifies the minimum and maximum latitudes and longitudes in the dataset.
    """
    # overwrite hard boundaries on longitude and latitude if not bounded in config.yaml
    if not options["bound_lon"]:
        grid_params["min_lon"] = 180
        grid_params["max_lon"] = -180
    if not options["bound_lat"]:
        grid_params["min_lat"] = 90
        grid_params["max_lat"] = -90

    # holds all ais data, one dataframe per csv file
    ais_data = []

    # get data from all csv files
    for csv_file in csv_files:
        # reads in the raw data with columns and number of rows specified in config.yaml
        nrows = options["max_rows"] if options["limit_rows"] else None
        usecols = ["MMSI", "LON", "LAT", "BaseDateTime"]
        ais_df = pd.read_csv(csv_file, usecols=usecols, nrows=nrows)
        ais_df = ais_df[usecols]

        # interprets raw time entries as datetime objects and drops original column
        ais_df["TIME"] = pd.to_datetime(
            ais_df["BaseDateTime"], format="%Y-%m-%dT%H:%M:%S"
        )
        ais_df.drop(columns="BaseDateTime", inplace=True)

        # keeps only rows in boundaries if specified
        if options["bound_lon"]:
            ais_df = ais_df.loc[
                (ais_df["LON"] >= grid_params["min_lon"])
                & (ais_df["LON"] <= grid_params["max_lon"])
            ]
        if options["bound_lat"]:
            ais_df = ais_df.loc[
                (ais_df["LAT"] >= grid_params["min_lat"])
                & (ais_df["LAT"] <= grid_params["max_lat"])
            ]

        # infers grid boundaries if no boundaries are specified
        if (
            not options["bound_lon"]
            and ais_df["LON"].min() < grid_params["min_lon"]
        ):
            grid_params["min_lon"] = ais_df["LON"].min()
        if (
            not options["bound_lon"]
            and ais_df["LON"].max() > grid_params["max_lon"]
        ):
            grid_params["max_lon"] = ais_df["LON"].max()
        if (
            not options["bound_lat"]
            and ais_df["LAT"].min() < grid_params["min_lat"]
        ):
            grid_params["min_lat"] = ais_df["LAT"].min()
        if (
            not options["bound_lat"]
            and ais_df["LAT"].max() > grid_params["max_lat"]
        ):
            grid_params["max_lat"] = ais_df["LAT"].max()

        # appends current dataframe to list of all dataframes
        ais_data.append(ais_df)

    # merges dataframes from all csvs
    trajectories = pd.concat(ais_data, axis=0, ignore_index=True)

    # rounds inferred grid boundaries to nearest degree to provide some padding to each boundary
    if not options["bound_lon"]:
        grid_params["min_lon"] = float(math.floor(grid_params["min_lon"]))
        grid_params["max_lon"] = float(math.ceil(grid_params["max_lon"]))
    if not options["bound_lat"]:
        grid_params["min_lat"] = float(math.floor(grid_params["min_lat"]))
        grid_params["max_lat"] = float(math.ceil(grid_params["max_lat"]))

    # number of columns in the resulting grid
    grid_params["num_cols"] = math.ceil(
        (grid_params["max_lon"] - grid_params["min_lon"])
        / grid_params["grid_len"]
    )

    return trajectories, grid_params

def write_data(trajectories, options, directories, grid_params):
    """Writes all trajectories to an output csv file using a discretized state and action grid.

    Uses the trajectories variable to look at each id-state-action-state transition to discretize all states and to
    interpolate actions if specified. These discretized states with interpolated or arbitrary actions are then written
    to the output csv specified by ``options['out_dir'] + options['out_file']``. Each trajectory is also sorted by
    its timestamp.

    The trajectories have their ids aliased by a counter variable that only increments whenever a trajectory is
    going to appear in the final csv. Self-transitions are discarded, and because of the huge grid size, most
    trajectories will be discarded since they will never transition between grid squares.

    Args:
        trajectories (pandas.DataFrame): All data entries with columns ``['MMSI', 'LON', 'LAT', 'TIME']``.
        options (dict): The script options specified in the ``config_file``.
        directories (dict): The input and output paths and files specified in the  ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.
    """
    # sorts based on MMSI, then sorts by timestamps within MMSI groups, drops the time column
    trajectories.sort_values(["MMSI", "TIME"], inplace=True)
    trajectories.drop(columns="TIME", inplace=True)

    # creates a new column of discretized states based on coordinate pairs
    trajectories["STATE"] = get_state(
        trajectories["LON"].values, trajectories["LAT"].values, grid_params
    )

    # looks at state differences within MMSI trajectories and only keeps the states with nonzero differences
    # trajectories with only one state are kept because they will have a first row with 'nan' for diff
    non_self_transitions = (
        trajectories["STATE"].groupby(trajectories["MMSI"]).diff().ne(0)
    )
    trajectories = trajectories.loc[non_self_transitions]

    # rounds latitude and longitude to specified precision
    trajectories = trajectories.round(
        {"LON": options["prec_coords"], "LAT": options["prec_coords"]}
    )

    # drops the trajectories with fewer states than ``options['min_states']``
    traj_lengths = trajectories["MMSI"].value_counts()
    traj_keep = traj_lengths[
        traj_lengths > options["min_states"] - 1
    ].index.values
    trajectories = trajectories.loc[trajectories["MMSI"].isin(traj_keep)]

    # aliases the MMSI column to ascending integers to enumerate trajectories and make easier to read
    alias = {
        mmsi: ind for ind, mmsi in enumerate(trajectories["MMSI"].unique())
    }
    trajectories["MMSI"] = trajectories["MMSI"].map(alias)

    # resets index now that manipulation of this dataframe has finished
    trajectories.reset_index(drop=True, inplace=True)

    # creates a series of stacked dataframes, each dataframe representing an interpolated state transition
    sas = trajectories.groupby("MMSI").apply(
        lambda x: get_action(x, options, grid_params)
    )
    if isinstance(
        sas, pd.DataFrame
    ):  # becomes a DataFrame when every trajectory has only one sas triplet
        sas = sas[0]

    # merge Series of dictionaries
    ids = []
    prevs = []
    acts = []
    curs = []
    lons = []
    lats = []
    for traj in sas:
        ids += traj["ID"]
        prevs += traj["PREV"]
        acts += traj["ACT"]
        curs += traj["CUR"]
        if options["append_coords"]:
            lons += traj["LON"]
            lats += traj["LAT"]

    # prepare final dictionary with built lists and proper heading name
    sas_data = {
        "sequence_id": ids,
        "from_state_id": prevs,
        "action_id": acts,
        "to_state_id": curs,
    }
    if options["append_coords"]:
        sas_data["lon"] = lons
        sas_data["lat"] = lats

    # writes new dataframe to final csv
    sas = pd.DataFrame(sas_data)
    sas.to_csv(
        directories["out_dir_path"] + directories["out_dir_file"], index=False
    )

def get_bounds(zone):
    """Helper function to get longitude boundaries corresponding to zone.

    Calculates the minimum and maximum longitudes corresponding to an integer zone
    representing a Universal Transverse Mercator coordinate system zone. Each zone is
    6 degrees wide, dividing the Earth into 60 zones, starting with zone 1 at 180 deg W. This function
    also wraps the zone with a modulo operator, so zone -1 would map to zone 58.

    Args:
        zone (int): The Universal Transverse Mercator coordinate system zone.

    Returns:
        tuple: The minimum and maximum longitudes of the zone passed in.
    """
    min_lon = (
        6.0 * ((zone - 1) % 60)
    ) - 180.0  # counts 6 degrees per zone, offset by -180

    return min_lon, (min_lon + 6.0)

def get_meta_data(file_name):
    """Helper function to retrieve a given file name's year, month, and zone.

    Takes a string file_name formatted as ``'AIS_yyyy_mm_Zone##.csv'`` and returns the numerical
    values of ``yyyy, mm, ##`` corresponding to year, month, and zone number as a tuple.

    Args:
        file_name (str): The file name to be parsed in format ``'AIS_yyyy_mm_Zone##.csv'``.

    Returns:
        tuple: The year, month, and zone corresponding to the filename passed in.
    """
    meta_file_data = file_name.split(
        "_"
    )  # splits csv file on '_' character, which separates relevant file info
    year = int(meta_file_data[-3])  # third to last element of file is the year
    month = int(
        meta_file_data[-2]
    )  # second to last element of file is the month

    # get zone number for csv file being read
    ending_raw = meta_file_data[
        -1
    ]  # gets last part of the file, with format "ZoneXX.csv"
    ending_data = ending_raw.split(
        "."
    )  # splits last part of file on '.' character
    zone_raw = ending_data[0]  # gets "ZoneXX" string
    zone_data = zone_raw[
        -2:
    ]  # gets last 2 characters of "ZoneXX" string - will be the zone number
    zone = int(zone_data)

    return year, month, zone

def get_state(cur_lon, cur_lat, grid_params):
    """Discretizes a coordinate pair into its state space representation in a Euclidean grid.

    Takes in a coordinate pair ``cur_lon``, ``cur_lat`` and grid parameters to calculate the integer state representing
    the given coordinate pair. This coordinate grid is always row-major. ``(min_lon, min_lat)`` represent the
    bottom-left corner of the grid.

    Example:
        A 3 x 4 grid would have the following state enumeration pattern::

            8 9 10 11
            4 5 6 7
            0 1 2 3

        With each grid square's area bounded in the following way::

            (min_lon, min_lat + grid_len)              (min_lon + grid_len, min_lat + grid_len)
                      |
             (min_lon, min_lat) ---------------------- (min_lon + grid_len, min_lon)

        In this example, the bottom left of state 0's boundaries would be the point ``min_lon, min_lat``, and the total
        area mapping to state 0 would be the square with ``min_lon, min_lat`` as the bottom left corner and each side of
        the square with length ``grid_len``. The inclusive parts of the square's boundaries mapping to zero are solid
        lines.

    Args:
        cur_lon (float): The longitude of the data point.
        cur_lat (float): The latitude of the data point.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        int: A state corresponding to the discretized representation of ``cur_lon``, ``cur_lat``.
    """
    # normalize lat and lon to the minimum values
    norm_lon = cur_lon - grid_params["min_lon"]
    norm_lat = cur_lat - grid_params["min_lat"]

    # find the row and column position based on grid_len
    col = norm_lon // grid_params["grid_len"]
    row = norm_lat // grid_params["grid_len"]

    # find total state based on num_cols in final grid
    return (row * grid_params["num_cols"] + col).astype(int)

def get_action(traj, options, grid_params):
    """Wrapper function for other ``get_action`` functions.

    Calls the correct ``get_action`` variant based on the options input and returns the resulting output with
    interpolated actions for all entries in the series.

    Args:
        traj (pandas.DataFrame): A pandas DataFrame with all the states encountered in a trajectory with their
            respective coordinates.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        dict: The sequence of state-action-state triplets for the passed in trajectory.
    """
    # retrieves trajectory data
    traj_num = traj.name
    states = traj["STATE"]
    last_state = states.iloc[-1]
    lon = traj["LON"]
    lat = traj["LAT"]

    # prepares a dictionary of state transitions to be fed row-by-row as a DataFrame to the interpolation functions
    data = {
        "ID": [traj_num] * (len(states) - 1),
        "PREV": states.iloc[:-1].values,
        "CUR": states.iloc[1:].values,
    }

    # if specified, appends the original entry coordinates (not discretized) for each 'PREV' entry
    if options["append_coords"]:
        data["LON"] = lon.iloc[:-1].values
        data["LAT"] = lat.iloc[:-1].values

    # formats the final data dictionary as a DataFrame
    traj_df = pd.DataFrame(data)

    # selects specified interpolation function and applies it row-wise to ``traj_df``
    if not options["interp_actions"]:
        traj_df = traj_df.apply(
            lambda x: get_action_arb(x, options, grid_params), axis=1
        )
    else:
        if options["allow_diag"]:
            traj_df = traj_df.apply(
                lambda x: get_action_interp_with_diag(x, options, grid_params),
                axis=1,
            )
        else:
            traj_df = traj_df.apply(
                lambda x: get_action_interp_reg(x, options, grid_params),
                axis=1,
            )

    # merges the dictionary series
    states_out = []
    acts_out = []
    lon_out = []
    lat_out = []
    for traj in traj_df:
        states_out += traj["PREV"]
        acts_out += traj["ACT"]
        if options["append_coords"]:
            lon_out += traj["LON"]
            lat_out += traj["LAT"]
    states_out.append(last_state)

    # appends the final state to each trajectory as its own row to allow for easier plotting of trajectories
    if options["append_coords"]:
        states_out.append(-1)
        acts_out.append(-1)
        lon_out.append(lon.iloc[-1])
        lat_out.append(lat.iloc[-1])

    # instantiates final dataframe-ready dictionary with state-action-state triplets
    data_out = {
        "ID": [traj_num] * len(acts_out),
        "PREV": states_out[:-1],
        "ACT": acts_out,
        "CUR": states_out[1:],
    }

    # adds coordinate fields for final output if specified in options
    if options["append_coords"]:
        data_out["LON"] = lon_out
        data_out["LAT"] = lat_out

    return data_out

def get_action_arb(row, options, grid_params):
    """Calculates an arbitrary action from the previous state to current state relative to the previous state.

    First, the relative offset between the current and previous state in rows and columns is calculated.
    The action is then calculated according to a spiral rule beginning with the previous state, so self-transitions
    are defined as ``0`` as an initial condition. Spiral inspired by the polar function ``r = theta``.

    Example:
        For example, if ``prev_state = 5``, ``cur_state = 7``, and ``num_cols = 4``, then our state grid is populated
        as follows::

            8  9 10 11
            4  p  6  c
            0  1  2  3

        Where p represents the location of the previous state, and c represents the location of the current state.
        Then the current state's position relative to the previous state is ``rel_row = 0``, ``rel_col = 2``. Our action
        spiral then looks like this::

            15 14 13 12 11      15 14 13 12 11
            16  4  3  2 10      16  4  3  2 10
            17  5  0  1  9  ->  17  5  p  1  c
            18  6  7  8 24      18  6  7  8 24
            19 20 21 22 23      19 20 21 22 23

        Thus, this algorithm will return ``9`` as the action.

    Args:
        row (pandas.Series): One row of the DataFrame the function is applied to, containing the trajectory number,
            previous state, current state, longitude, and latitude.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        dict: State-action-state triplets that interpolate between ``prev_state`` and ``cur_state``.
    """
    # retrieves transition data
    traj_num = row["ID"].astype(int)
    prev_state = row["PREV"].astype(int)
    cur_state = row["CUR"].astype(int)
    num_cols = grid_params["num_cols"]

    # gets row, column decomposition for previous and current states
    prev_row = prev_state // num_cols
    prev_col = prev_state % num_cols
    cur_row = cur_state // num_cols
    cur_col = cur_state % num_cols
    # calculates current state's position relative to previous state
    rel_row = cur_row - prev_row
    rel_col = cur_col - prev_col

    # simple routine to calculate a spiral set of actions
    # the sequence defined by layer corresponds to the total number of grid squares in each spiral layer
    action_num = x = y = i = 0
    layer = (2 * i + 1) ** 2  # sets breakpoint for when to increment i
    while not (x == rel_col and y == rel_row):
        if action_num == layer - 1:
            i += 1  # move to next spiral
            x = i
            layer = (2 * i + 1) ** 2  # calculate breakpoint for next spiral
        elif (
            x == i and y < i
        ):  # traverses from beginning of layer to top right corner
            y += 1
        elif x > -i and y == i:  # traverses from top right to top left corner
            x -= 1
        elif (
            x == -i and y > -i
        ):  # traverses from top left to bottom left corner
            y -= 1
        elif (
            x < i and y == -i
        ):  # traverses from bottom left to bottom right corner
            x += 1
        elif (
            x == i and y < 0
        ):  # traverses from bottom left corner to end of layer
            y += 1
        action_num += 1

    # prepares final data dictionary to build DataFrame
    out_data = {
        "ID": [traj_num],
        "PREV": [prev_state],
        "ACT": [action_num],
        "CUR": [cur_state],
    }

    # overwrites the coordinates of the first state in interpolated transitions to be original raw values
    if options["append_coords"]:
        out_data["LON"] = [row["LON"]]
        out_data["LAT"] = [row["LAT"]]

    return out_data

def get_action_interp_with_diag(row, options, grid_params):
    """Calculates the actions taken from the previous state to reach the current state, interpolating if necessary.

    First, the relative offset between the current and previous state in rows and columns is calculated.
    Then the sign of ``rel_row`` and ``rel_col`` are then used to iteratively describe a sequence of actions
    from the previous state to current state, breaking up state transitions with multiple actions if
    the states are not adjacent (including diagonals, resulting in 9 possible actions). This interpolation
    assumes a deterministic system.

    Example:
        For example, if ``prev_state = 5``, ``cur_state = 7``, and ``num_cols = 4``, then our state grid is populated
        as follows::

            8  9 10 11
            4  p  6  c
            0  1  2  3

        Output snippet::

            pd.DataFrame({})

        Where p represents the location of the previous state, and c represents the location of the current state.
        Then the current state's position relative to the previous state is ``rel_row = 0``, ``rel_col = 2``. Our
        action spiral then looks like this::

            4  3  2      4  3  2
            5  0  1  ->  5  p  1  c
            7  8  9      6  7  8

        Output snippet::

            pd.DataFrame({
                            'ID': [traj_num, ],
                            'PREV': [prev_state, ],
                            'ACT': [1, ],
                            'CUR': [prev_state + 1, ]
                        })

        Because the current state is not adjacent (including diagonals), we interpolate by taking the action that
        brings us closest to the current state: action ``1``, resulting in a new action spiral and a new previous
        state::

            4  3  2      4  3  2
            5  0  1  ->  5  p  c
            7  8  9      6  7  8

        Final output::

            pd.DataFrame({
                            'ID': [traj_num] * 2,
                            'PREV': [prev_state, prev_state + 1],
                            'ACT': [1, 1],
                            'CUR': [prev_state + 1, cur_state]
                        })

        Now, our new previous state is adjacent to the current state, so we can take action ``1``, which updates our
        previous state to exactly match the current state, so the algorithm terminates and returns the list of
        state-action-state transitions.

    Args:
        row (pandas.Series): One row of the DataFrame the function is applied to, containing the trajectory number,
            previous state, current state, longitude, and latitude.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        dict: State-action-state triplets that interpolate between ``prev_state`` and ``cur_state``.
    """
    # retrieves transition data
    traj_num = row["ID"].astype(int)
    prev_state = row["PREV"].astype(int)
    cur_state = row["CUR"].astype(int)
    num_cols = grid_params["num_cols"]

    # instantiate lists to hold column values for final DataFrame output
    prevs = []
    acts = []
    curs = []
    lons = []
    lats = []

    # gets row, column decomposition for previous and current states
    prev_row = prev_state // num_cols
    prev_col = prev_state % num_cols
    cur_row = cur_state // num_cols
    cur_col = cur_state % num_cols
    # calculates current state's position relative to previous state
    rel_row = cur_row - prev_row
    rel_col = cur_col - prev_col

    # write output rows until rel_row and rel_col are both zero
    # out_rows = []
    while not (rel_row == 0 and rel_col == 0):
        # selects action to minimize rel_row and rel_col
        action = -1
        if rel_row > 0 and rel_col > 0:
            action = 2
        elif rel_row > 0 and rel_col == 0:
            action = 3
        elif rel_row > 0 and rel_col < 0:
            action = 4
        elif rel_row == 0 and rel_col > 0:
            action = 1
        elif rel_row == 0 and rel_col < 0:
            action = 5
        elif rel_row < 0 and rel_col > 0:
            action = 8
        elif rel_row < 0 and rel_col == 0:
            action = 7
        elif rel_row < 0 and rel_col < 0:
            action = 6

        # moves rel_row and rel_col in the opposite directions of their signs
        row_diff = -np.sign(rel_row)
        col_diff = -np.sign(rel_col)

        # updates states and relative row, column based on action selected
        rel_row += row_diff
        rel_col += col_diff
        temp_row = prev_row - row_diff
        temp_col = prev_col - col_diff
        temp_state = temp_row * num_cols + temp_col
        prev_state = prev_row * num_cols + prev_col

        # records an interpolated state-action-state transition
        prevs.append(prev_state)
        acts.append(action)
        curs.append(temp_state)

        # gets the coordinates of the interpolated state - will be the coordinates of the middle of the state
        if options["append_coords"]:
            lon, lat = state_to_coord(prev_state, options, grid_params)
            lons.append(lon)
            lats.append(lat)

        prev_row = temp_row
        prev_col = temp_col

    # prepares final data dictionary to build DataFrame
    out_data = {
        "ID": [traj_num] * len(prevs),
        "PREV": prevs,
        "ACT": acts,
        "CUR": curs,
    }

    # overwrites the coordinates of the first state in interpolated transitions to be original raw values
    if options["append_coords"]:
        lons[0] = row["LON"]
        lats[0] = row["LAT"]
        out_data["LON"] = lons
        out_data["LAT"] = lats

    return out_data

def get_action_interp_reg(row, options, grid_params):
    """Calculates the actions taken from the previous state to reach the current state, interpolating if necessary.

    First, the relative offset between the current and previous state in rows and columns is calculated.
    Then the sign of ``rel_row`` and ``rel_col`` are then used to iteratively describe a sequence of actions
    from the previous state to current state, breaking up state transitions with multiple actions if
    the states are not adjacent (only actions are right, left, up, down, and none). This interpolation
    assumes a deterministic system.

    Example:
        For example, if ``prev_state = 5``, ``cur_state = 7``, and ``num_cols = 4``, then our state grid is populated
        as follows::

            8  9 10 11
            4  p  6  c
            0  1  2  3

        Output snippet::

            pd.DataFrame({})

        Where p represents the location of the previous state, and c represents the location of the current state.
        Then the current state's position relative to the previous state is ``rel_row = 0``, ``rel_col = 2``. Our action
        spiral then looks like this::

               2            2
            3  0  1  ->  3  p  1  c
               4            4

        Output snippet::

            output: pd.DataFrame({
                                    'ID': [traj_num, ],
                                    'PREV': [prev_state, ],
                                    'ACT': [1, ],
                                    'CUR': [prev_state + 1, ]
                                })

        Because the current state is not adjacent, we interpolate by taking the action that brings us closest to
        the current state: action ``1``, resulting in a new action spiral and a new previous state::

               2            1
            3  0  1  ->  2  p  c
               4            4

        Final output::

            pd.DataFrame({
                            'ID': [traj_num] * 2,
                            'PREV': [prev_state, prev_state + 1],
                            'ACT': [1, 1],
                            'CUR': [prev_state + 1, cur_state]
                        })

        Now, our new previous state is adjacent to the current state, so we can take action ``1``, which updates our
        previous state to exactly match the current state, so the algorithm terminates and returns the list of
        state-action-state transitions.

    Args:
        row (pandas.Series): One row of the DataFrame the function is applied to, containing the trajectory number,
            previous state, current state, longitude, and latitude.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        dict: State-action-state triplets that interpolate between ``prev_state`` and ``cur_state``.
    """
    # retrieves transition data
    traj_num = row["ID"].astype(int)
    prev_state = row["PREV"].astype(int)
    cur_state = row["CUR"].astype(int)
    num_cols = grid_params["num_cols"]

    # instantiate lists to hold column values for final DataFrame output
    prevs = []
    acts = []
    curs = []
    lons = []
    lats = []

    # gets row, column decomposition for previous and current states
    prev_row = prev_state // num_cols
    prev_col = prev_state % num_cols
    cur_row = cur_state // num_cols
    cur_col = cur_state % num_cols
    # calculates current state's position relative to previous state
    rel_row = cur_row - prev_row
    rel_col = cur_col - prev_col

    # write output rows until rel_row and rel_col are both zero
    while not (rel_row == 0 and rel_col == 0):
        # selects action to reduce the largest of rel_row and rel_col
        action = -1
        if rel_row > 0 and rel_col > 0:
            action = 2 if rel_row > rel_col else 1
        elif rel_row > 0 and rel_col == 0:
            action = 2
        elif rel_row > 0 and rel_col < 0:
            action = 2 if rel_row > -rel_col else 3
        elif rel_row == 0 and rel_col > 0:
            action = 1
        elif rel_row == 0 and rel_col < 0:
            action = 3
        elif rel_row < 0 and rel_col > 0:
            action = 4 if -rel_row > rel_col else 1
        elif rel_row < 0 and rel_col == 0:
            action = 4
        elif rel_row < 0 and rel_col < 0:
            action = 4 if -rel_row > -rel_col else 3

        # moves rel_row and rel_col in the opposite directions of their signs
        row_diff = -np.sign(rel_row) if (action == 2 or action == 4) else 0
        col_diff = -np.sign(rel_col) if (action == 1 or action == 3) else 0

        # updates states and relative row, column based on action selected
        rel_row += row_diff
        rel_col += col_diff
        temp_row = prev_row - row_diff
        temp_col = prev_col - col_diff
        temp_state = temp_row * num_cols + temp_col
        prev_state = prev_row * num_cols + prev_col

        # records an interpolated state-action-state transition
        prevs.append(prev_state)
        acts.append(action)
        curs.append(temp_state)

        # gets the coordinates of the interpolated state - will be the coordinates of the middle of the state
        if options["append_coords"]:
            lon, lat = state_to_coord(prev_state, options, grid_params)
            lons.append(lon)
            lats.append(lat)

        prev_row = temp_row
        prev_col = temp_col

    # prepares final data dictionary to build DataFrame
    out_data = {
        "ID": [traj_num] * len(acts),
        "PREV": prevs,
        "ACT": acts,
        "CUR": curs,
    }

    # overwrites the coordinates of the first state in interpolated transitions to be original raw values
    if options["append_coords"]:
        lons[0] = row["LON"]
        lats[0] = row["LAT"]
        out_data["LON"] = lons
        out_data["LAT"] = lats

    return out_data

def state_to_coord(state, options, grid_params):
    """Inverse function for ``get_state``.

    Calculates the coordinates of the middle of the passed in state in the specified grid passed in.

    Args:
        state (int): The discretized grid square returned by ``get_state``.
        options (dict): The script options specified in the ``config_file``.
        grid_params (dict): The grid parameters specified in the ``config_file``.

    Returns:
        tuple: The longitude and latitude representing the middle of the state passed in.
    """
    # calculates the integer state's row and column representation in the grid
    state_col = state % grid_params["num_cols"]
    state_row = state // grid_params["num_cols"]

    # calculates the latitude and longitude corresponding to the middle of the grid square
    state_lon = round(
        grid_params["min_lon"] + grid_params["grid_len"] * (state_col + 0.5),
        options["prec_coords"],
    )
    state_lat = round(
        grid_params["min_lat"] + grid_params["grid_len"] * (state_row + 0.5),
        options["prec_coords"],
    )

    return state_lon, state_lat

if __name__ == "__main__":
    main()

max_stars_repo_name: jpalat/AIS-Vessel-Data-Pipeline

max_issues_repo_name: adewin/AIS-Vessel-Data-Pipeline

max_forks_repo_name: adewin/AIS-Vessel-Data-Pipeline


text:
from itertools import chain
from math import acos
from numpy import array
from scipy.optimize import minimize

from model.Bonus import Bonus
from model.Building import Building
from model.Faction import Faction
from model.Minion import Minion
from model.MinionType import MinionType
from model.StatusType import StatusType
from model.Tree import Tree
from model.Wizard import Wizard

from strategy_common import Point, Line, Circle

BONUS_WEIGHT = 1
DIRECTION_WEIGHT = 1
PROJECTILE_WEIGHT = 1
TARGET_DISTANCE_WEIGHT = 1
UNIT_WEIGHT = 1

def get_optimal_position(target, me: Wizard, buildings, minions, wizards, trees, projectiles, bonuses,
                         orc_woodcutter_attack_range, fetish_blowdart_attack_range, magic_missile_direct_damage,
                         magic_missile_radius, dart_radius, map_size, shielded_direct_damage_absorption_factor,
                         empowered_damage_factor, max_distance=None, max_iterations=None, penalties=None):
    my_position = Point(me.x, me.y)

    def filter_friends(units):
        return (v for v in units if v.faction == me.faction)

    def filter_enemies(units):
        return (v for v in units if is_enemy(v, me.faction))

    def filter_neutral(units):
        return (v for v in units if v.faction == Faction.NEUTRAL)

    def filter_max_distance(units):
        if max_distance is None:
            return units
        else:
            return (v for v in units if my_position.distance(v.position) <= max_distance)

    enemy_buildings = tuple(filter_max_distance(filter_enemies(buildings)))
    enemy_minions = tuple(filter_max_distance(filter_enemies(minions)))
    enemy_wizards = tuple(filter_max_distance(filter_enemies(wizards)))
    bonuses = tuple(filter_max_distance(bonuses))
    if not enemy_buildings and not enemy_minions and not enemy_wizards and not bonuses:
        factor = 2 if me.mean_speed.norm() < 1 else 1.3
        trees = tuple(v for v in trees if my_position.distance(v.position) < factor * me.radius + v.radius)
        if trees:
            return min(trees, key=lambda v: my_position.distance(v.position)).position
        neutral_minions = tuple(v for v in filter_neutral(minions)
                                if my_position.distance(v.position) < factor * me.radius + v.radius)
        if neutral_minions:
            return min(neutral_minions, key=lambda v: my_position.distance(v.position)).position
        return my_position
    get_damage = make_get_damage(magic_missile_direct_damage, empowered_damage_factor)
    get_attack_range = make_get_attack_range(
        orc_woodcutter_attack_range=orc_woodcutter_attack_range,
        fetish_blowdart_attack_range=fetish_blowdart_attack_range,
        magic_missile_radius=magic_missile_radius,
        dart_radius=dart_radius,
        my_radius=me.radius,
    )
    other_wizards = tuple(v for v in wizards if v.id != me.id)
    units = tuple(v for v in chain(buildings, trees, minions, other_wizards) if id(v) != id(target))
    friends_wizards = tuple(filter_friends(chain(other_wizards)))
    friends_units = tuple(chain(friends_wizards, filter_friends(chain(minions, buildings))))
    enemies_units = tuple(filter_enemies(chain(buildings, minions, other_wizards)))
    target_position = target.position if target else None
    damage_factor = 1 - next((shielded_direct_damage_absorption_factor * bool(v.remaining_duration_ticks)
                              for v in me.statuses if v.type == StatusType.SHIELDED), 0)

    def position_penalty(values):
        position = Point(values[0], values[1])

        def get_unit_damage(unit):
            attack_range = get_attack_range(unit)
            distance = position.distance(unit.position)
            return get_damage(unit) * (1 if attack_range >= distance
                                       else distance_penalty(distance - attack_range, 0.5 * attack_range))

        sum_damage = sum(damage_factor * get_unit_damage(v) for v in enemies_units)

        def unit_intersection_penalty(unit):
            if isinstance(unit, Tree):
                return distance_penalty(position.distance(unit.position), 2 * me.radius + 2 * unit.radius)
            elif isinstance(unit, Minion) and unit.faction == Faction.NEUTRAL:
                return distance_penalty(position.distance(unit.position), 2 * me.radius + 2 * unit.radius)
            else:
                return distance_penalty(position.distance(unit.position), 1.1 * me.radius + unit.radius)

        def unit_danger_penalty(unit):
            if not is_enemy(unit, me.faction):
                return 0
            distance_to_unit = position.distance(unit.position)
            if isinstance(unit, Minion) and friends_units:
                nearest_friend = min(friends_units, key=lambda v: v.position.distance(unit.position))
                if nearest_friend.position.distance(unit.position) < distance_to_unit:
                    return 0
            add_damage = damage_factor * (get_damage(unit) - get_unit_damage(unit))
            safe_distance = max(me.cast_range + magic_missile_radius,
                                (get_attack_range(unit) + 2 * me.radius) *
                                min(1, 2 * (sum_damage + add_damage) / me.life))
            return distance_penalty(distance_to_unit, safe_distance)

        def target_distance_penalty():
            if not target:
                return 0
            if isinstance(target, Bonus):
                return bonus_penalty(target)
            else:
                max_cast_range = me.cast_range + magic_missile_radius
                distance = position.distance(target_position)
                if distance <= max_cast_range:
                    d_penalty = 0
                else:
                    safe_distance = max(max_cast_range, map_size)
                    d_penalty = 1 - distance_penalty(distance - max_cast_range, safe_distance)
                return max(unit_danger_penalty(target), d_penalty)

        def friend_unit_intersection_penalty(friend):
            circle = Circle(friend.position, friend.radius)
            intersection = circle.has_intersection_with_moving_circle(
                Circle(position, magic_missile_radius), target_position)
            if not intersection:
                return 0
            target_to_friend = friend.position - target_position
            tangent_cos = (friend.radius + magic_missile_radius) / target_position.distance(friend.position)
            tangent_angle = acos(min(1, max(-1, tangent_cos)))
            tangent1_direction = target_to_friend.rotate(tangent_angle)
            tangent2_direction = target_to_friend.rotate(-tangent_angle)
            tangent1 = target_position + tangent1_direction
            tangent2 = target_position + tangent2_direction
            tangent1_distance = Line(target_position, tangent1).distance(position)
            tangent2_distance = Line(target_position, tangent2).distance(position)
            max_distance = (tangent1_distance + tangent2_distance) * 0.5
            distance_to_tangent = min(tangent1_distance, tangent2_distance)
            return distance_to_tangent / max_distance

        def friend_units_intersections_penalties():
            return (friend_unit_intersection_penalty(v) for v in friends_wizards)

        def direction_penalty():
            return max(friend_units_intersections_penalties()) if target_position and friends_wizards else 0

        def bonus_penalty(bonus):
            return 1 - distance_penalty(position.distance(bonus.position),
                                        my_position.distance(bonus.position) + me.radius)

        def units_penalties():
            for unit in units:
                yield max(unit_intersection_penalty(unit), unit_danger_penalty(unit))

        def bonuses_penalties():
            return (bonus_penalty(v) for v in bonuses)

        def projectile_penalty(projectile):
            projectile_speed = projectile.mean_speed
            safe_distance = 2 * me.radius + projectile.radius
            distance_to = Line(projectile.position, projectile.position + projectile_speed).distance(position)
            return distance_penalty(distance_to, safe_distance)

        def projectiles_penalties():
            return (projectile_penalty(v) for v in projectiles)

        def borders_penalty():
            return (0 if me.radius < position.x < map_size - me.radius
                    and me.radius < position.y < map_size - me.radius else 1)

        penalties_count = (
            + len(units) * UNIT_WEIGHT
            + len(bonuses) * BONUS_WEIGHT
            + len(projectiles) * PROJECTILE_WEIGHT
            + (DIRECTION_WEIGHT if target_position and friends_wizards else 0)
            + (TARGET_DISTANCE_WEIGHT if target_position else 0)
        )
        penalty = max(borders_penalty() * penalties_count, (
            + sum(units_penalties()) * UNIT_WEIGHT
            + sum(bonuses_penalties()) * BONUS_WEIGHT
            + sum(projectiles_penalties()) * PROJECTILE_WEIGHT
            + direction_penalty() * DIRECTION_WEIGHT
            + target_distance_penalty() * TARGET_DISTANCE_WEIGHT
        ))
        if penalties is not None:
            penalties.append((position, penalty))
        return penalty

    from_my_position = minimize(position_penalty, array([my_position.x, my_position.y]),
                                method='Nelder-Mead', options=dict(maxiter=max_iterations))
    return Point(from_my_position.x[0], from_my_position.x[1])

def get_target(me: Wizard, buildings, minions, wizards, trees, bonuses, magic_missile_direct_damage,
               empowered_damage_factor, staff_range, max_distance=None):
    my_position = Point(me.x, me.y)

    def filter_enemies(units):
        return (v for v in units if is_enemy(v, me.faction))

    def filter_neutral(units):
        return (v for v in units if v.faction == Faction.NEUTRAL)

    def filter_max_distance(units):
        if max_distance is None:
            return units
        else:
            return (v for v in units if my_position.distance(v.position) <= max_distance)

    enemy_buildings = tuple(filter_max_distance(filter_enemies(buildings)))
    enemy_minions = tuple(filter_max_distance(filter_enemies(minions)))
    enemy_wizards = tuple(filter_max_distance(filter_enemies(wizards)))
    bonuses = tuple(filter_max_distance(bonuses))
    if not enemy_buildings and not enemy_minions and not enemy_wizards and not bonuses:
        factor = 2 if me.mean_speed.norm() < 1 else 1.3
        trees = tuple(v for v in trees if my_position.distance(v.position) < factor * me.radius + v.radius)
        if trees:
            return min(trees, key=lambda v: my_position.distance(v.position))
        neutral_minions = tuple(v for v in filter_neutral(minions)
                                if my_position.distance(v.position) < factor * me.radius + v.radius)
        if neutral_minions:
            return min(neutral_minions, key=lambda v: my_position.distance(v.position))
        return None
    get_damage = make_get_damage(magic_missile_direct_damage, empowered_damage_factor)

    def target_penalty(unit):
        distance = unit.position.distance(my_position)
        return distance * unit.life / get_damage(me) if distance <= 2 * unit.radius + me.vision_range else distance

    def is_in_staff_range(unit):
        return my_position.distance(unit.position) <= staff_range + unit.radius

    def get_optimal_in_range(units, is_in_range):
        in_range = tuple(v for v in units if is_in_range(v))
        return min(in_range, key=target_penalty) if in_range else None

    if bonuses:
        return min(bonuses, key=lambda v: my_position.distance(v.position))
    elif enemy_wizards or enemy_minions or enemy_buildings:
        target = get_optimal_in_range(enemy_wizards, is_in_staff_range)
        if target is None:
            target = get_optimal_in_range(enemy_minions, is_in_staff_range)
        if target is None:
            if enemy_wizards:
                target = min(enemy_wizards, key=target_penalty)
            elif enemy_minions:
                target = min(enemy_minions, key=target_penalty)
            elif enemy_buildings:
                target = min(enemy_buildings, key=target_penalty)
        return target

NOT_ENEMIES = frozenset((Faction.NEUTRAL, Faction.OTHER))

def is_enemy(unit, my_faction):
    return unit.faction != my_faction and unit.faction not in NOT_ENEMIES

def make_get_attack_range(orc_woodcutter_attack_range, fetish_blowdart_attack_range, magic_missile_radius,
                          dart_radius, my_radius):
    minions = {
        MinionType.ORC_WOODCUTTER: orc_woodcutter_attack_range,
        MinionType.FETISH_BLOWDART: fetish_blowdart_attack_range,
    }

    def impl(unit):
        if isinstance(unit, Building):
            return unit.attack_range + my_radius
        if isinstance(unit, Minion):
            if unit.type == MinionType.ORC_WOODCUTTER:
                return minions[unit.type] + my_radius
            elif unit.type == MinionType.FETISH_BLOWDART:
                return minions[unit.type] + my_radius + dart_radius
        if isinstance(unit, Wizard):
            return unit.cast_range + my_radius + magic_missile_radius

    return impl

def make_get_damage(magic_missile_direct_damage, empowered_damage_factor):

    def impl(unit):
        factor = 1 + next((empowered_damage_factor for v in unit.statuses if v.type == StatusType.EMPOWERED), 0)
        if isinstance(unit, (Building, Minion)):
            return factor * unit.damage
        if isinstance(unit, Wizard):
            return factor * magic_missile_direct_damage

    return impl

def distance_penalty(value, safe_distance):
    return max(0, (safe_distance - value) / safe_distance)

ext: py

max_stars_repo_name: elsid/CodeWizards

max_issues_repo_name: elsid/CodeWizards

max_forks_repo_name: elsid/CodeWizards

reason: from numpy,from scipy


text:
# pip3 install opencv-contrib-python

import numpy as np
import cv2 as cv

img1 = cv.imread("box.png")  # queryImage
img2 = cv.imread("box_in_scene.png")  # trainImage

# Initiate BRISK detector
detector = cv.BRISK_create()

# Create Matcher window
cv.namedWindow("DAISY BF Matcher", cv.WINDOW_NORMAL)

def f(x):
    return

# BoostDesc descriptor types
norm_types = {
    0: 100,  # NRM_NONE
    1: 101,  # NRM_PARTIAL
    2: 102,  # NRM_FULL
    3: 103,  # NRM_SIFT
}

# Set DAISY descriptor parameter callbacks
cv.createTrackbar("Radius", "DAISY BF Matcher", 15, 50, f)
cv.createTrackbar("Q Radius", "DAISY BF Matcher", 2, 50, f)
cv.createTrackbar("Q Theta", "DAISY BF Matcher", 8, 50, f)
cv.createTrackbar("Q Hist", "DAISY BF Matcher", 8, 50, f)
cv.createTrackbar("Norm", "DAISY BF Matcher", 0, len(norm_types) - 1, f)
cv.createTrackbar("Interpolation", "DAISY BF Matcher", 1, 1, f)
cv.createTrackbar("Use Orientation", "DAISY BF Matcher", 0, 1, f)

while True:

    # find the keypoints with BRISK
    kp1 = detector.detect(img1, None)
    kp2 = detector.detect(img2, None)

    # Read matcher parameters
    current_radius = cv.getTrackbarPos("Radius", "DAISY BF Matcher")
    current_q_radius = cv.getTrackbarPos("Q Radius", "DAISY BF Matcher") + 1
    current_q_theta = cv.getTrackbarPos("Q Theta", "DAISY BF Matcher")
    current_q_hist = cv.getTrackbarPos("Q Hist", "DAISY BF Matcher")
    current_norm = cv.getTrackbarPos("Norm", "DAISY BF Matcher")
    current_interpolation = cv.getTrackbarPos("Interpolation", "DAISY BF Matcher")
    current_orientation = cv.getTrackbarPos("Use Orientation", "DAISY BF Matcher")

    # Initiate DAISY
    descriptor = cv.xfeatures2d.DAISY_create(
        radius=current_radius,
        q_radius=current_q_radius,
        q_theta=current_q_theta,
        q_hist=current_q_hist,
        norm=norm_types[current_norm],
        interpolation=current_interpolation,
        use_orientation=current_orientation,
    )

    # find the descriptors with DAISY
    kp1, des1 = descriptor.compute(img1, kp1)
    kp2, des2 = descriptor.compute(img2, kp2)

    # create BFMatcher object
    bf = cv.BFMatcher(cv.NORM_L2, crossCheck=True)

    # Match descriptors.
    matches = bf.match(des1, des2)

    # Sort them in the order of their distance.
    matches = sorted(matches, key=lambda x: x.distance)

    # Flags:
    # cv.DRAW_MATCHES_FLAGS_DEFAULT
    # cv.DRAW_MATCHES_FLAGS_DRAW_OVER_OUTIMG
    # cv.DRAW_MATCHES_FLAGS_NOT_DRAW_SINGLE_POINTS
    # cv.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS
    img3 = cv.drawMatches(
        img1,
        kp1,
        img2,
        kp2,
        matches,
        None,
        flags=cv.DRAW_MATCHES_FLAGS_NOT_DRAW_SINGLE_POINTS,
        matchColor=(0, 255, 0),
    )

    # Draw matches
    cv.imshow("DAISY BF Matcher", img3)

    # Calculate homography
    # Consider point filtering
    obj = []
    scene = []
    for match in matches:
        obj.append(kp1[match.queryIdx].pt)
        scene.append(kp2[match.trainIdx].pt)

    # Calculate homography: Inliers and outliers
    # RANSAC, LMEDS, RHO
    H, _ = cv.findHomography(np.array(obj), np.array(scene), cv.RANSAC)

    if H is not None:
        # Frame of the object image
        obj_points = np.array(
            [
                [0, 0],
                [img1.shape[1], 0],
                [img1.shape[1], img1.shape[0]],
                [0, img1.shape[0]],
            ],
            dtype=np.float,
        )

        # Check the sanity of the transformation
        warped_points = cv.perspectiveTransform(np.array([obj_points]), H)

        warped_image = np.copy(img2)
        cv.drawContours(
            warped_image, np.array([warped_points]).astype(np.int32), 0, (0, 0, 255)
        )

        cv.namedWindow("Warped Object", cv.WINDOW_NORMAL)
        cv.imshow("Warped Object", warped_image)
    else:
        print("Error calculating perspective transformation")

    if cv.waitKey(10) & 0xFF == 27:
        break

cv.destroyAllWindows()

max_stars_repo_name: vladiant/ComputerVisionUtils

max_issues_repo_name: vladiant/ComputerVisionUtils

max_forks_repo_name: vladiant/ComputerVisionUtils

reason: import numpy


text:
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Log generator for tensor data."""
import time

from operator import mul
from functools import reduce
import numpy as np
from mindinsight.datavisual.proto_files import mindinsight_anf_ir_pb2 as anf_ir_pb2
from mindinsight.datavisual.proto_files import mindinsight_summary_pb2 as summary_pb2

from .log_generator import LogGenerator

class TensorLogGenerator(LogGenerator):
    """
    Log generator for tensor data.

    This is a log generator writing tensor data. User can use it to generate fake
    summary logs about tensor.
    """

    def generate_event(self, values):
        """
        Method for generating tensor event.

        Args:
            values (dict): A dict contains:
                {
                    wall_time (float): Timestamp.
                    step (int): Train step.
                    value (float): Tensor value.
                    tag (str): Tag name.
                }

       Returns:
            summary_pb2.Event.

        """
        tensor_event = summary_pb2.Event()
        tensor_event.wall_time = values.get('wall_time')
        tensor_event.step = values.get('step')

        value = tensor_event.summary.value.add()
        value.tag = values.get('tag')
        tensor = values.get('value')

        value.tensor.dims[:] = tensor.get('dims')
        value.tensor.data_type = tensor.get('data_type')
        value.tensor.float_data[:] = tensor.get('float_data')

        return tensor_event

    def generate_log(self, file_path, steps_list, tag_name):
        """
        Generate log for external calls.

        Args:
            file_path (str): Path to write logs.
            steps_list (list): A list consists of step.
            tag_name (str): Tag name.

        Returns:
            list[dict], generated tensor metadata.
            list, generated tensors.

        """
        tensor_metadata = []
        tensor_values = dict()
        for step in steps_list:
            tensor = dict()

            wall_time = time.time()
            tensor.update({'wall_time': wall_time})
            tensor.update({'step': step})
            tensor.update({'tag': tag_name})
            dims = list(np.random.randint(1, 10, 4))
            mul_value = reduce(mul, dims)
            tensor.update({'value': {
                "dims": dims,
                "data_type": anf_ir_pb2.DataType.DT_FLOAT32,
                "float_data": np.random.randn(mul_value)
            }})
            tensor_metadata.append(tensor)
            tensor_values.update({step: tensor})

            self._write_log_one_step(file_path, tensor)

        return tensor_metadata, tensor_values

if __name__ == "__main__":
    tensor_log_generator = TensorLogGenerator()
    test_file_name = '%s.%s.%s' % ('tensor', 'summary', str(time.time()))
    test_steps = [1, 3, 5]
    test_tag = "test_tensor_tag_name"
    tensor_log_generator.generate_log(test_file_name, test_steps, test_tag)

max_stars_repo_name: fapbatista/mindinsight

max_issues_repo_name: fapbatista/mindinsight

max_forks_repo_name: fapbatista/mindinsight


text:
import numpy as np
from keras.layers import Input
from keras.models import Model

from categorical_embedder.embedders.core.seq2seq.Seq2SeqTrainerABC import Seq2SeqTrainerABC
from categorical_embedder.embedders.core.seq2seq.EncoderDecoder import EncoderDecoder
from categorical_embedder.processors.HolisticCategoricalProcessor import HolisticCategoricalProcessor

class Trainer(Seq2SeqTrainerABC):
    def __init__(self, num_of_categories, uniques, outer_generator, generator_info, model_info, save_info):
        super().__init__(HolisticCategoricalProcessor(uniques, 0, np.int), outer_generator, generator_info, model_info, save_info)

        self._num_of_categories = num_of_categories
        self._vocab_size = self._preprocessor.get_vocab_size()

    def _get_model(self):
        encoder_input = Input((self._num_of_categories, self._vocab_size), dtype="int64")
        decoder_input = Input((self._num_of_categories + 1, self._vocab_size), dtype="int64")

        enc_dec = EncoderDecoder(self._vocab_size, self._model_info["embedding_length"], "softmax" if self._model_info["softmax_on"] else "linear", self._model_info["encoder_info"], self._model_info["decoder_info"])

        decoder_output, _ = enc_dec([encoder_input, decoder_input])
        model = Model(inputs=[encoder_input, decoder_input], outputs=decoder_output)

        self._main_model_artifacts["custom_objects_info"]["layer_info"] = ["EncoderDecoder", "StackedRecurrentEncoder", "StackedRecurrentDecoder"]
        self._embedder_artifacts["custom_objects_info"]["layer_info"] = ["StackedRecurrentEncoder"]

        return model

max_stars_repo_name: erelcan/categorical-embedder

max_issues_repo_name: erelcan/categorical-embedder

max_forks_repo_name: erelcan/categorical-embedder


text:
import faiss
import torch
import numpy as np
from tqdm import tqdm
from torch.utils.data import DataLoader

def compute_features(geoloc_dataset, model, global_features_dim, num_workers=4,
                     eval_batch_size=32, recall_values=[1, 5, 10, 20]):
    """Compute the features of all images within the geoloc_dataset.

    Parameters
    ----------
    geoloc_dataset : dataset_geoloc.GeolocDataset,  which contains the images (queries and gallery).
    model : network.Network.
    global_features_dim : int, dimension of the features (e.g. 256 for AlexNet with GeM).
    num_workers : int.
    eval_batch_size : int.
    recall_values : list of int, recalls to compute (e.g. R@1, R@5...).

    Returns
    -------
    recalls : np.array of int, containing R@1, R@5, r@10, r@20.
    recalls_pretty_str : str, pretty-printed recalls.
    predictions : np.array of int, containing the first 20 predictions for each query,
        with shape [queries_num, 20].
    correct_bool_mat : np.array of int, with same dimension of predictions,
        indicates of the prediction is correct or wrong. Its values are only [0, 1].
    distances : np.array of float, with same dimension of predictions,
        indicates the distance in features space from the query to its prediction.
    """
    test_dataloader = DataLoader(dataset=geoloc_dataset, num_workers=num_workers,
                                 batch_size=eval_batch_size, pin_memory=True)
    model = model.eval()
    with torch.no_grad():
        gallery_features = np.empty((len(geoloc_dataset), global_features_dim), dtype="float32")
        for inputs, indices in tqdm(test_dataloader, desc=f"Comp feats {geoloc_dataset}", ncols=120):
            B, C, H, W = inputs.shape
            inputs = inputs.cuda()
            # Compute outputs using global features (e.g. GeM, NetVLAD...)
            output = model("features_extractor", [inputs, "global"])
            output = output.reshape(B, global_features_dim)
            gallery_features[indices.detach().numpy(), :] = output.detach().cpu().numpy()
    query_features = gallery_features[geoloc_dataset.gallery_num:]
    gallery_features = gallery_features[:geoloc_dataset.gallery_num]
    faiss_index = faiss.IndexFlatL2(global_features_dim)
    faiss_index.add(gallery_features)

    max_recall_value = max(recall_values)  # Usually it's 20
    distances, predictions = faiss_index.search(query_features, max_recall_value)
    ground_truths = geoloc_dataset.get_positives()

    recalls, recalls_str = compute_recalls(predictions, ground_truths, geoloc_dataset, recall_values)

    correct_bool_mat = np.zeros((geoloc_dataset.queries_num, max_recall_value), dtype=np.int)
    for query_index in range(geoloc_dataset.queries_num):
        positives = set(ground_truths[query_index].tolist())
        for pred_index in range(max_recall_value):
            pred = predictions[query_index, pred_index]
            if pred in positives:
                correct_bool_mat[query_index, pred_index] = 1
    return recalls, recalls_str, predictions, correct_bool_mat, distances

def compute_recalls(predictions, ground_truths, test_dataset, recall_values=[1, 5, 10, 20]):
    """Computes the recalls.

    Parameters
    ----------
    predictions : np.array of int, containing the first 20 predictions for each query,
        with shape [queries_num, 20].
    ground_truths : list of lists of int, containing for each query the list of its positives.
        It's a list of lists because each query has different amount of positives.
    test_dataset : dataset_geoloc.GeolocDataset.
    recall_values : list of int, recalls to compute (e.g. R@1, R@5...).

    Returns
    -------
    recalls : np.array of int, containing R@1, R@5, r@10, r@20.
    recalls_pretty_str : str, pretty-printed recalls.
    """

    recalls = np.zeros(len(recall_values))
    for query_index, pred in enumerate(predictions):
        for i, n in enumerate(recall_values):
            if np.any(np.in1d(pred[:n], ground_truths[query_index])):
                recalls[i:] += 1
                break
    recalls = recalls / test_dataset.queries_num * 100
    recalls_pretty_str = ", ".join([f"R@{val}: {rec:.1f}" for val, rec in zip(recall_values, recalls)])
    return recalls, recalls_pretty_str


text:
import numpy as np
import glob
import os
import random
import tqdm
import torch
from torchvision import transforms, utils
from PIL import Image
from skvideo.io import vwrite
import imageio
import pickle
from copy import deepcopy
from my_utils.visualize_flame_overlay import OverLayViz
from my_utils.ringnet_overlay.util import tensor_vis_landmarks
from dataset_loaders import fast_image_reshape
import constants
from my_utils.flm_dynamic_fit_overlay import camera_dynamic
import torchvision

# # useful constant
SHAPE_IDS = constants.INDICES['SHAPE']
EXP_IDS = constants.INDICES['EXP']
POSE_IDS = constants.INDICES['POSE']
TRANS_IDS = constants.INDICES['TRANS']
# CAM_IDS = constants.INDICES['CAM']
# JAW_ROT_IDS = INDICES['JAW_ROT']
# GLOBAL_ROT_IDS = INDICES['GLOBAL_ROT']
# ROT_JAW_CAM = INDICES['ROT_JAW_CAM']

flame_params_path = '../../New_Flame.txt._is_scratch'

def get_original_images(flame_path):
    image_path = flame_path.replace('/New_flame/', '/Cropped_frames/')
    return Image.open(image_path.replace('npy', 'png'))

def read_flame_param(flame_path, data_type = 'VOXCELEB'):
    '''
        Return flame parameters: [0:100 - Shape, 100:150 - Expression, 150:156 - Pose, 156:159 - Camera]
    '''
    if data_type == 'VOXCELEB':
        flame_path_old = flame_path.replace('/New_flame/', '/Flame_params/')

        flame_new = np.load(flame_path)
        flame_old = np.load(flame_path_old)

        params = flame_old[:]
        # params[0:100] = np.array(shape_params[id])
        params[100:150] = flame_new[:50]
        params[153:156] = flame_new[53:56]

        # TODO: SEND 156 dimensional param vector where params[:100] corresponds to shape
        params = params[:156]

    elif data_type == 'FFHQ':
        with open(flame_path, 'rb') as param_file:
            flame_param = pickle.load(param_file, encoding='latin1')

        params = np.hstack((flame_param['betas'][:100], flame_param['betas'][300:350], flame_param['pose'][0:3],
                                 flame_param['pose'][6:9]))

    return params

def make_flame_shape_consitant(flame_params):
     flame_median = np.median(flame_params[:,:100], axis=0)
     flame_params[:, :100] = flame_median
     return flame_params

def save_images_to_gif(images, save_path):
    '''
        Input: images: Batch_Size x 3 x H x W [Torch.Tensor]
    '''
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    images = images.clamp(min=0, max=1.0)
    imgs = images.detach().cpu().numpy().transpose((0, 2, 3, 1)) * 255
    imgs = imgs.astype('uint8')

    imageio.imwrite(save_path.replace('.gif', '.png'), imgs[0])
    imageio.mimsave(save_path, imgs)

    vwrite(save_path.replace('.gif', '.mp4'), imgs, outputdict={'-r': '30'})

def save_images_with_flame_to_gif(images, save_path):
    pass
    '''
        Input: images: Batch_Size x 3 x H x W [Torch.Tensor]
    '''
    # batch_size = images.shape[0]
    # frames = []
    #
    # for img in images:
    #     if isinstance(img, Image.Image):
    #         new_frame = img
    #     else:
    #         new_frame = transforms.ToPILImage()((img.cpu().data + 1)/2)
    #     frames.append(new_frame)
    #
    # frames[0].save(save_path, format='GIF', append_images=frames[1:], save_all=True, duration=41, loop=0)

def generate_from_flame_sequence(generator, gen_in, pose, step, alpha, noise, input_indices):
    '''
        Input: flame_params: Batch_Size x Flame_Code_Size
        Output: fake_images: Batch_Size x 3 x H x W
    '''
    max_batch_szie = 32
    generator.eval()

    images = []
    with torch.no_grad():
        for batch_idx in range(0, gen_in.shape[0], max_batch_szie):
            gen_in_this_batch = gen_in[batch_idx:batch_idx + max_batch_szie]

            if type(gen_in) is np.ndarray:
                gen_in_this_batch = torch.from_numpy(gen_in_this_batch).cuda()

            if pose is not None and type(pose) is np.ndarray:
                pose_this_batch = torch.from_numpy(pose[batch_idx:batch_idx + max_batch_szie]).cuda()
            else:
                pose_this_batch = None

            if input_indices.shape[0] == 1:
                identities = input_indices.repeat((gen_in_this_batch.shape[0],))
            else:
                identities = input_indices[batch_idx:batch_idx + max_batch_szie]

            fake_images = generator(gen_in_this_batch, pose_this_batch, step=step, alpha=alpha, noise=noise,
                                    input_indices=identities)
            images.append(fake_images[-1].cpu())

    if len(images) > 1:
        images = [torch.cat(images, dim=0)]

    return images

def generate_gif_from_vgg_params(generator, flame_params_path, normalization_file_path, log_dir, step, noise, n_rows=5,
                                 n_cols=4, seed=0, n_frames=8, dump_original=False):
    random.seed(seed)
    os.makedirs(log_dir, exist_ok=True)

    grid_size = n_rows * n_cols
    with open(flame_params_path, 'r') as f:
        flame_param_files = f.readlines()

    if normalization_file_path is not None:
        normalization_data = np.load(normalization_file_path)
        flame_mean = normalization_data['mean'][0]
        flame_std = normalization_data['std'][0]

    # pick N=grid_size different random sequence
    cell_ids = random.sample(flame_param_files, 50)

    for i, cell in enumerate(tqdm.tqdm(cell_ids)):
        cell_dir = os.path.dirname(cell)
        flame_files = sorted(glob.glob(os.path.join(cell_dir, '*.npy')))[:n_frames]
        assert len(flame_files) >= n_frames, 'Number of video frames are' + str(len(flame_files))

        flame_params = np.array(list(map(read_flame_param, flame_files)))
        flame_params = make_flame_shape_consitant(flame_params)

        # Keep only strongly meaningful dimensions
        # flame_params[:, :100] = 0
        # flame_params[:, 100:150] = 0
        # flame_params[153:156] = 0
        # flame_params[153:156]

        # Normalization
        if normalization_file_path is not None:
            flame_params = (flame_params - flame_mean)/flame_std

        fake_images = generate_from_flame_sequence(generator, flame_params, step, noise)[-1]

        save_path = os.path.join(log_dir, str(i).zfill(4)+'.gif')
        save_images_to_gif(fake_images, save_path)

        if dump_original:
            path_dirs = save_path.split('/')
            save_path_original = path_dirs[:-1] + ['original'] + [str(i).zfill(4)+'.gif']
            original_images = list(map(get_original_images, flame_files))
            save_images_to_gif(original_images, os.path.join(*save_path_original))

def get_gif_from_list_of_params(generator, flame_params, step, alpha, noise, overlay_landmarks, flame_std, flame_mean,
                                overlay_visualizer, rendered_flame_as_condition, use_posed_constant_input,
                                normal_maps_as_cond, camera_params):
    # cam_t = np.array([0., 0., 2.5])
    # camera_params = camera_dynamic((224, 224), cam_t)
    if overlay_visualizer is None:
        overlay_visualizer = OverLayViz()

    fixed_embeddings = torch.ones(flame_params.shape[0], dtype=torch.long, device='cuda')*13
    # print(generator.module.get_embddings()[fixed_embeddings])
    flame_params_unnorm = flame_params * flame_std + flame_mean

    flame_params_unnorm = torch.from_numpy(flame_params_unnorm).cuda()
    normal_map_img, _, _, _, rend_imgs = \
        overlay_visualizer.get_rendered_mesh(flame_params=(flame_params_unnorm[:, SHAPE_IDS[0]:SHAPE_IDS[1]],
                                                           flame_params_unnorm[:, EXP_IDS[0]:EXP_IDS[1]],
                                                           flame_params_unnorm[:, POSE_IDS[0]:POSE_IDS[1]],
                                                           flame_params_unnorm[:, TRANS_IDS[0]:TRANS_IDS[1]]),
                                             camera_params=camera_params)
    rend_imgs = (rend_imgs/127.0 - 1)

    if use_posed_constant_input:
        pose = flame_params[:, constants.get_idx_list('GLOBAL_ROT')]
    else:
        pose = None

    if rendered_flame_as_condition:
        gen_in = rend_imgs
    else:
        gen_in = flame_params

    if normal_maps_as_cond:
        gen_in = torch.cat((rend_imgs, normal_map_img), dim=1)

    fake_images = generate_from_flame_sequence(generator, gen_in, pose, step, alpha, noise,
                                               input_indices=fixed_embeddings)[-1]

    fake_images = overlay_visualizer.range_normalize_images(fast_image_reshape(fake_images,
                                                                               height_out=rend_imgs.shape[2],
                                                                               width_out=rend_imgs.shape[3]))
    if rendered_flame_as_condition:
        fake_images = torch.cat([fake_images.cpu(), (rend_imgs.cpu() + 1)/2], dim=-1)

    if normal_maps_as_cond:
        fake_images = torch.cat([fake_images.cpu(), (normal_map_img.cpu() + 1) / 2], dim=-1)

    return fake_images

def interpolate_FFHQ(generator, flame_param_paths, normalization_file_path, log_dir, step, alpha, noise, camera_params,
                     seed=0, n_frames=8, flame_params=None, overlay_landmarks=False, flame_mean=None, flame_std=None,
                     rendered_flame_as_condition=False, overlay_visualizer=None, use_posed_constant_input=False,
                     normal_maps_as_cond=False):
    '''
        cell1: Vary shape
        cell2: Vary expression
        cell3: Vary Global Pose
        cell4: Vary Jaw Pose
    '''

    random.seed(seed)
    os.makedirs(log_dir, exist_ok=True)

    if flame_params is None:
        flame_param1 = read_flame_param(flame_param_paths[0], data_type='FFHQ')
        flame_param2 = read_flame_param(flame_param_paths[1], data_type='FFHQ')

        if normalization_file_path is not None:
            normalization_data = np.load(normalization_file_path)
            flame_mean = normalization_data['mean'][0]
            flame_std = normalization_data['std'][0]
            flame_param1 = (flame_param1 - flame_mean)/flame_std
            flame_param2 = (flame_param2 - flame_mean)/flame_std
    else:
        flame_param1, flame_param2 = (flame_params - flame_mean)/flame_std

    # cell_ids = [SHAPE_IDS, EXP_IDS, GLOBAL_ROT_IDS, JAW_ROT_IDS, ROT_JAW_CAM]
    # cell_names = ['SHAPE', 'EXP', 'GLOBAL_ROT', 'JAW_ROT', ('GLOBAL_ROT', 'JAW_ROT'), ('JAW_ROT', 'EXP'), 'TRANS',
    #               ('GLOBAL_ROT', 'TRANS'), ('SHAPE', 'GLOBAL_ROT', 'TRANS')]
    cell_names = [('JAW_ROT', 'EXP'),  ('GLOBAL_ROT', 'TRANS')]

    start_end_images = get_gif_from_list_of_params(generator, np.vstack((flame_param1, flame_param2)), step, alpha,
                                                   noise, overlay_landmarks, flame_std, flame_mean, overlay_visualizer,
                                                   rendered_flame_as_condition, use_posed_constant_input,
                                                   normal_maps_as_cond, camera_params)
    torchvision.utils.save_image(
        start_end_images,
        os.path.join(log_dir, 'start_end_overlayed.png'),
        nrow=2,
        normalize=True,
        range=(0, 1))

    for cell_name in tqdm.tqdm(cell_names):
        # cell_dir = os.path.dirname(cell)
        # flame_files = sorted(glob.glob(os.path.join(cell_dir, '*.npy')))[:n_frames]
        # assert len(flame_files) >= n_frames, 'Number of video frames are' + str(len(flame_files))
        cell_idxs = constants.get_idx_list(cell_name)
        flame_params = np.expand_dims(deepcopy(flame_param1), axis=0).repeat(n_frames, axis=0).astype('float32')
        # import ipdb; ipdb.set_trace()
        flame_params[:, cell_idxs] = np.linspace(flame_param1[cell_idxs], flame_param2[cell_idxs], n_frames,
                                                    endpoint=True)

        # # hold some flame prams constant
        # flame_params[:, 5:100] = flame_params[0, 5:100]
        # flame_params[:, 105:150] = flame_params[0, 105:150]
        # flame_params *= 0

        fake_images = get_gif_from_list_of_params(generator, flame_params, step, alpha, noise, overlay_landmarks,
                                                  flame_std, flame_mean, overlay_visualizer,
                                                  rendered_flame_as_condition, use_posed_constant_input,
                                                  normal_maps_as_cond, camera_params)
        # fake_images = fake_images.clamp_(0, 1)

        if type(cell_name) == tuple:
            save_name = ''
            for name_frg in cell_name:
                save_name += (name_frg + '_')
        else:
            save_name = cell_name

        save_path = os.path.join(log_dir, save_name + '.gif')
        # save_images_with_flame_to_gif(fake_images, flame, save_path)
        save_images_to_gif(fake_images, save_path)

max_stars_repo_name: qiaone/GIF

max_issues_repo_name: qiaone/GIF

max_forks_repo_name: qiaone/GIF


text:
import os
from fnmatch import fnmatch
from typing import Iterator, List, Optional, Sequence, Tuple, Union

import numpy as np
from typing_extensions import Literal

from . import lib
from .otf import TemporaryOTF
from .util import PathOrArray, _kwargs_for, imread

def rl_cleanup():
    """Release GPU buffer and cleanup after deconvolution

    Call this before program quits to release global GPUBuffer d_interpOTF.

    - Resets any bleach corrections
    - Removes OTF from GPU buffer
    - Destroys cuFFT plan
    - Releases GPU buffers
    """
    return lib.RL_cleanup()

def rl_init(
    rawdata_shape: Tuple[int, int, int],
    otfpath: str,
    dzdata: float = 0.5,
    dxdata: float = 0.1,
    dzpsf: float = 0.1,
    dxpsf: float = 0.1,
    deskew: float = 0,
    rotate: float = 0,
    width: int = 0,
):
    """Initialize GPU for deconvolution.

    Prepares cuFFT plan for deconvolution with a given data shape and OTF.
    Must be used prior to :func:`pycudadecon.rl_decon`

    Parameters
    ----------
    rawdata_shape : Tuple[int, int, int]
        3-tuple of data shape
    otfpath : str
        Path to OTF TIF
    dzdata : float, optional
        Z-step size of data, by default 0.5
    dxdata : float, optional
        XY pixel size of data, by default 0.1
    dzpsf : float, optional
        Z-step size of the OTF, by default 0.1
    dxpsf : float, optional
        XY pixel size of the OTF, by default 0.1
    deskew : float, optional
        Deskew angle. If not 0.0 then deskewing will be performed before
        deconvolution, by default 0
    rotate : float, optional
        Rotation angle; if not 0.0 then rotation will be performed around Y
        axis after deconvolution, by default 0
    width : int, optional
        If deskewed, the output image's width, by default 0 (do not crop)

    Examples
    --------
    >>> rl_init(im.shape, otfpath)
    >>> decon_result = rl_decon(im)
    >>> rl_cleanup()

    """
    nz, ny, nx = rawdata_shape
    lib.RL_interface_init(
        nx,
        ny,
        nz,
        dxdata,
        dzdata,
        dxpsf,
        dzpsf,
        deskew,
        rotate,
        width,
        otfpath.encode(),
    )

def rl_decon(
    im: np.ndarray,
    background: Union[int, Literal["auto"]] = 80,
    n_iters: int = 10,
    shift: int = 0,
    save_deskewed: bool = False,
    output_shape: Optional[Tuple[int, int, int]] = None,
    napodize: int = 15,
    nz_blend: int = 0,
    pad_val: float = 0.0,
    dup_rev_z: bool = False,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """Perform Richardson Lucy Deconvolution.

    Performs actual deconvolution. GPU must first be initialized with
    :func:`pycudadecon.rl_init`

    Parameters
    ----------
    im : np.ndarray
        3D image volume to deconvolve
    background : int or 'auto'
        User-supplied background to subtract. If 'auto', the median value of the
        last Z plane will be used as background. by default 80
    n_iters : int, optional
        Number of iterations, by default 10
    shift : int, optional
        If deskewed, the output image's extra shift in X (positive->left),
        by default 0
    save_deskewed : bool, optional
        Save deskewed raw data as well as deconvolution result, by default False
    output_shape : tuple of int, optional
        Specify the output shape after deskewing.  Usually this is unnecessary and
        can be autodetected.  Mostly intended for use within a
        :class:`pycudadecon.RLContext` context, by default None
    napodize : int, optional
        Number of pixels to soften edge with, by default 15
    nz_blend : int, optional
        Number of top and bottom sections to blend in to reduce axial ringing,
        by default 0
    pad_val : float, optional
        Value with which to pad image when deskewing, by default 0.0
    dup_rev_z : bool, optional
        Duplicate reversed stack prior to decon to reduce axial ringing,
        by default False

    Returns
    -------
    np.ndarray or 2-tuple of np.ndarray
        The deconvolved result.  If `save_deskewed` is `True`, returns
        `(decon_result, deskew_result)`

    Raises
    ------
    ValueError
        If im.ndim is not 3, or `output_shape` is provided but not length 3
    """

    if im.ndim != 3:
        raise ValueError("Only 3D arrays supported")

    nz, ny, nx = im.shape
    if output_shape is None:
        output_shape = (lib.get_output_nz(), lib.get_output_ny(), lib.get_output_nx())
    elif len(output_shape) != 3:
        raise ValueError("Decon output shape must have length==3")

    decon_result = np.empty(tuple(output_shape), dtype=np.float32)

    if save_deskewed:
        deskew_result = np.empty_like(decon_result)
    else:
        deskew_result = np.empty(1, dtype=np.float32)

    # must be 16 bit going in
    if not np.issubdtype(im.dtype, np.uint16):
        im = im.astype(np.uint16)

    if isinstance(background, str) and background == "auto":
        background = np.median(im[-1])

    rescale = False  # not sure if this works yet...

    if not im.flags["C_CONTIGUOUS"]:
        im = np.ascontiguousarray(im)
    lib.RL_interface(
        im,
        nx,
        ny,
        nz,
        decon_result,
        deskew_result,
        background,
        rescale,
        save_deskewed,
        n_iters,
        shift,
        napodize,
        nz_blend,
        pad_val,
        dup_rev_z,
    )

    if save_deskewed:
        return decon_result, deskew_result
    else:
        return decon_result

def quickDecon(image: np.ndarray, otfpath: str, **kwargs):
    """Perform deconvolution of `image` with otf at `otfpath`.

    Not currently used...
    """
    rl_init(image.shape, otfpath, **_kwargs_for(rl_init, kwargs))
    result = rl_decon(image, **_kwargs_for(rl_decon, kwargs))
    lib.RL_cleanup()
    return result

class RLContext:
    """Context manager to setup the GPU for RL decon

    Takes care of handing the OTF to the GPU, preparing a cuFFT plane,
    and cleaning up after decon.  Internally, this calls :func:`rl_init`,
    stores the shape of the expected output volume after any deskew/decon,
    then calls :func:`rl_cleanup` when exiting the context.

    For parameters, see :func:`rl_init`.

    Examples
    --------
    >>> with RLContext(data.shape, otfpath, dz) as ctx:
    ...     result = rl_decon(data, ctx.out_shape)
    """

    def __init__(
        self,
        rawdata_shape: Tuple[int, int, int],
        otfpath: str,
        dzdata: float = 0.5,
        dxdata: float = 0.1,
        dzpsf: float = 0.1,
        dxpsf: float = 0.1,
        deskew: float = 0,
        rotate: float = 0,
        width: int = 0,
    ):
        self.kwargs = locals()
        self.kwargs.pop("self")
        self.out_shape: Optional[Tuple[int, int, int]] = None

    def __enter__(self):
        """Setup the context and return the ZYX shape of the output image"""
        rl_init(**self.kwargs)
        self.out_shape = (lib.get_output_nz(), lib.get_output_ny(), lib.get_output_nx())
        return self

    def __exit__(self, typ, val, traceback):
        # exit receives a tuple with any exceptions raised during processing
        # if __exit__ returns True, exceptions will be supressed
        lib.RL_cleanup()

# alias
rl_context = RLContext

def _yield_arrays(
    images: Union[PathOrArray, Sequence[PathOrArray]], fpattern="*.tif"
) -> Iterator[np.ndarray]:
    """Yield arrays from an array, path, or sequence of either.

    Parameters
    ----------
    images : Union[PathOrArray, Sequence[PathOrArray]]
        an array, path, or sequence of either
    fpattern : str, optional
        used to filter files in a directory, by default "*.tif"

    Yields
    -------
    Iterator[np.ndarray]
        Arrays (read from paths if necessary)

    Raises
    ------
    OSError
        If a directory is provided and no files match fpattern.
    """
    if isinstance(images, np.ndarray):
        yield images

    elif isinstance(images, str):
        if os.path.isfile(images):
            yield imread(images)

        elif os.path.isdir(images):
            imfiles = [f for f in os.listdir(images) if fnmatch(f, fpattern)]
            if not len(imfiles):
                raise OSError(
                    'No files matching pattern "{}" found in directory: {}'.format(
                        fpattern, images
                    )
                )
            for fpath in imfiles:
                yield imread(os.path.join(images, fpath))

    else:
        for item in images:
            yield from _yield_arrays(item)

def decon(
    images: Union[PathOrArray, Sequence[PathOrArray]],
    psf: PathOrArray,
    fpattern: str = "*.tif",
    **kwargs
) -> Union[np.ndarray, List[np.ndarray]]:
    """Deconvolve an image or images with a PSF or OTF file.

    If `images` is a directory, use the `fpattern` argument to select files
    by filename pattern.

    Parameters
    ----------
    images : str, np.ndarray, or sequence of either
        The array, filepath, directory, or list/tuple thereof to deconvolve
    psf : str or np.ndarray
        a filepath of a PSF or OTF file, or a 3D numpy PSF array.  Function will
        auto-detect whether the file is a 3D PSF or a filepath representing a 2D
        complex OTF.
    fpattern : str, optional
        Filepattern to use when a directory is provided in the `images` argument,
        by default `*.tif`
    ** kwargs
        All other kwargs must be valid for either :func:`rl_init` or :func:`rl_decon`.

    Returns
    -------
    np.ndarray or list of array
        The deconvolved image(s)

    Raises
    ------
    ValueError
        If save_deskewed is True and deskew is unset or 0
    IOError
        If a directory is provided as input and ``fpattern`` yields no files
    NotImplementedError
        If ``psf`` is provided as a complex, 2D numpy array (OTFs can only be
        provided as filenames created with :func:`pycudadecon.make_otf`)

    Examples
    --------
    deconvolve a 3D TIF volume with a 3D PSF volume (e.g. a single bead stack)

    >>> result = decon('/path/to/image.tif', '/path/to/psf.tif')

    deconvolve all TIF files in a specific directory that match a certain
    `filename pattern <https://docs.python.org/3.6/library/fnmatch.html>`_,
    (in this example, all TIFs with the string '560nm' in their name)

    >>> result = decon(
    ...     '/directory/with/images', '/path/to/psf.tif', fpattern='*560nm*.tif'
    ... )

    deconvolve a list of images, provided either as np.ndarrays, filepaths,
    or directories

    >>> imarray = tifffile.imread('some_other_image.tif')
    >>> inputs = ['/directory/with/images', '/path/to/image.tif', imarray]
    >>> result = decon(inputs, '/path/to/psf.tif', fpattern='*560nm*.tif')
    """
    if kwargs.get("save_deskewed"):
        if kwargs.get("deskew", 1) == 0:
            raise ValueError("Cannot use save_deskewed=True with deskew=0")
        if not kwargs.get("deskew"):
            raise ValueError("Must set deskew != 0 when using save_deskewed=True")

    init_kwargs = _kwargs_for(rl_init, kwargs)
    decon_kwargs = _kwargs_for(rl_decon, kwargs)

    out = []
    with TemporaryOTF(psf, **kwargs) as otf:
        arraygen = _yield_arrays(images, fpattern)
        # first, assume that all of the images are the same shape...
        # in which case we can prevent a lot of GPU IO
        # grab and store the shape of the first item in the generator
        next_im = next(arraygen)
        shp = next_im.shape

        with RLContext(shp, otf.path, **init_kwargs) as ctx:
            while True:
                out.append(
                    rl_decon(next_im, output_shape=ctx.out_shape, **decon_kwargs)
                )
                try:
                    next_im = next(arraygen)
                    # here we check to make sure that the images are still the same
                    # shape... if not, we'll continue below
                    if next_im.shape != shp:
                        break
                except StopIteration:
                    next_im = None
                    break

        # if we had a shape mismatch, there will still be images left to process
        # process them the slow way here...
        if next_im is not None:
            for imarray in [next_im, *arraygen]:
                with RLContext(imarray.shape, otf.path, **init_kwargs) as ctx:
                    out.append(
                        rl_decon(imarray, output_shape=ctx.out_shape, **decon_kwargs)
                    )

    if isinstance(images, (list, tuple)) and len(images) > 1:
        return out
    else:
        return out[0]

max_stars_repo_name: zbarry/pycudadecon

max_issues_repo_name: zbarry/pycudadecon

max_forks_repo_name: zbarry/pycudadecon


text:
import torch
import torch.nn as nn
from xmuda.models.DDR import Bottleneck3D
import numpy as np
from xmuda.models.resnet34_unet import UNetResNet34

class ProcessKitti(nn.Module):
    def __init__(self, feature, norm_layer, bn_momentum, dilations=[1, 2, 3]):
        super(Process, self).__init__()
        self.main = nn.Sequential(
            *[Bottleneck3D(feature,
                           feature // 4,
                           bn_momentum=bn_momentum,
                           norm_layer=norm_layer,
                           dilation=[i, i, i]) for i in dilations]
        )
    def forward(self, x):
        return self.main(x)

class Process(nn.Module):
    def __init__(self, feature, norm_layer, bn_momentum, dilations=[1, 2, 3]):
        super(Process, self).__init__()
        self.main = nn.Sequential(
            *[Bottleneck3D(feature,
                           feature // 4,
                           bn_momentum=bn_momentum,
                           norm_layer=norm_layer,
                           dilation=[i, i, i]) for i in dilations]
        )
    def forward(self, x):
        return self.main(x)

class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer, bn_momentum):
        super(Upsample, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose3d(in_channels, out_channels, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1),
            norm_layer(out_channels, momentum=bn_momentum),
            nn.ReLU()
        )

    def forward(self, x):
        return self.main(x)

class Downsample(nn.Module):
    def __init__(self, feature, norm_layer, bn_momentum, expansion=8):
        super(Downsample, self).__init__()
        self.main = Bottleneck3D(feature,
                                 feature // 4,
                                 bn_momentum=bn_momentum,
                                 expansion=expansion, stride=2,
                                 downsample=nn.Sequential(
                                     nn.AvgPool3d(kernel_size=2, stride=2),
                                     nn.Conv3d(feature, int(feature * expansion/4), kernel_size=1, stride=1, bias=False),
                                     norm_layer(int(feature * expansion/4), momentum=bn_momentum)
                                 ),
                                 norm_layer=norm_layer)

    def forward(self, x):
        return self.main(x)

max_stars_repo_name: anhquancao/xmuda-extend

max_issues_repo_name: anhquancao/xmuda-extend

max_forks_repo_name: anhquancao/xmuda-extend


text:
import cv2
import base64
from socketIO_client import SocketIO, BaseNamespace
import numpy as np
import time
from PIL import Image
from threading import Thread, ThreadError
import io

img_np = None
socketIO = SocketIO('http://192.168.0.1', 6020)
live_namespace = socketIO.define(BaseNamespace, '/live')

def receive_events_thread():
    socketIO.wait()
def on_camera_response(*args):
    print(args[0]['data'])

live_namespace.on('camera_update', on_camera_response)
receive_events_thread = Thread(target=receive_events_thread)
receive_events_thread.daemon = True
receive_events_thread.start()

cap = cv2.VideoCapture(1)
#cap.set(3, 1280)
#cap.set(4, 720)
while True:
    ret, img = cap.read()
    img_b = cv2.imencode('.jpg', img)[1].tobytes()
    base64_bytes = base64.b64encode(img_b)
    base64_string = base64_bytes.decode('utf-8')
    live_namespace.emit('livevideo',{'data':base64_string})
    time.sleep(10)

max_stars_repo_path: huseinhouse.com/Face Analysis with Elastic/client.py