text:
# -*- coding: utf-8 -*-

import numpy as np
import tomopy
from scipy.ndimage import zoom, gaussian_filter
# from scipy.stats import signaltonoise

import gc
import operator

def arrange_image(img1, img2, shift):

    shift = map(int, shift)
    new_shape = map(int, map(max, map(operator.add, img2.shape, shift), img1.shape))
    newimg = np.empty(new_shape)
    newimg[:, :] = np.NaN
    newimg[0:img1.shape[0], 0:img1.shape[1]] = img1
    newimg[shift[0]:shift[0] + img2.shape[0], shift[1]:shift[1] + img2.shape[1]] = img2
    gc.collect()

    return newimg

def trim_sinogram(data, center, x, y, diameter, fin_angle=180):
    """
    Adopted from Tomopy.
    Provide sinogram corresponding to a circular region of interest
    by trimming the complete sinogram of a compact object.

    Parameters
    ----------
    data : ndarray
        Input 3D data.

    center : float
        Rotation center location.

    x, y : int, int
        x and y coordinates in pixels (image center is (0, 0))

    diameter : float
        Diameter of the circle of the region of interest.

    Returns
    -------
    ndarray
        Output 3D data.

    """
    data = data.copy().astype('float32')
    dx, dy, dz = data.shape
    mask = np.zeros([dx, dz], dtype='bool')

    rad = np.sqrt(x * x + y * y)
    alpha = np.arctan2(x, y)
    l1 = center - diameter / 2
    l2 = center - diameter / 2 + rad

    roidata = np.ones((dx, dy, diameter), dtype='float32')

    delphi = np.pi * (fin_angle / 180) / dx
    for m in range(dx):

        # Calculate start end coordinates for each row of the sinogram.
        ind1 = np.ceil(np.cos(alpha - m * delphi) * (l2 - l1) + l1)
        ind2 = np.ceil(np.cos(alpha - m * delphi) * (l2 - l1) + l1 + diameter)

        # Make sure everythin is inside the frame.
        if ind1 < 0:
            ind1 = 0
        if ind1 > dz:
            ind1 = dz
        if ind2 < 0:
            ind2 = 0
        if ind2 > dz:
            ind2 = dz

        ind1, ind2 = map(int, [ind1, ind2])
        roidata[m, :, 0:(ind2 - ind1)] = data[m:m+1, :, ind1:ind2]
        mask[m, ind1:ind2] = True

    return roidata, mask

def normalize(img):

    return (img - img.min()) / (img.max() - img.min())

def ssim(x, y, terms='lcs', mask_ratio=None):

    ssim = 1
    mask = tomopy.misc.corr._get_mask(x.shape[0], x.shape[1], mask_ratio)
    x = x[mask]
    y = y[mask]
    for term in terms:
        if term == 'l':
            mu_x = np.mean(x)
            mu_y = np.mean(y)
            c1 = (0.01 * (np.max([x.max(), y.max()]) - np.min([x.min(), y.min()])))**2
            ssim *= ((2 * mu_x * mu_y + c1) / (mu_x**2 + mu_y**2 + c1))
        if term == 'c':
            sigma_x = np.sqrt(np.var(x))
            sigma_y = np.sqrt(np.var(y))
            c2 = (0.03 * (np.max([x.max(), y.max()]) - np.min([x.min(), y.min()])))**2
            ssim *= ((2 * sigma_x * sigma_y + c2) / (sigma_x**2 + sigma_y**2 + c2))**2
        if term == 's':
            t = np.vstack([x.flatten(), y.flatten()])
            t = np.cov(t)
            sigma_x = np.sqrt(t[0, 0])
            sigma_y = np.sqrt(t[1, 1])
            sigma_xy = t[0, 1]
            c3 = (0.03 * (np.max([x.max(), y.max()]) - np.min([x.min(), y.min()])))**2 / 2
            ssim *= ((sigma_xy + c3) / (sigma_x * sigma_y + c3))
    return ssim

def snr(img, ref, mask_ratio=None, ss_error=False):
    """
    Calculate the signal-to-noise ratio.
    :param img: image array
    :param ref: ground truth
    :return:
    """
    r = np.copy(ref)
    i = np.copy(img)
    if mask_ratio is None:
        rmean = r.mean()
        imean = i.mean()
        i = i + rmean - imean
    else:
        mask = tomopy.misc.corr._get_mask(img.shape[0], img.shape[1], mask_ratio)
        r = r[mask]
        i = i[mask]
        rmean = r.mean()
        imean = i.mean()
        i = i + rmean - imean
    if ss_error:
        return np.linalg.norm(r - i) ** 2
    else:
        return 10 * np.log10(np.linalg.norm(r) ** 2 / np.linalg.norm(r - i) ** 2)

def snr_intrinsic(img, mask_ratio=None):
    """
    Calculate the signal-to-noise ratio.
    :param img: image array
    :param ref: ground truth
    :return:
    """
    i = np.copy(img)
    if mask_ratio is not None:
        mask = tomopy.misc.corr._get_mask(img.shape[0], img.shape[1], mask_ratio)
        i = i[mask]
    snr = signaltonoise(i)
    return snr

def downsample_img(img, ds, axis=0):

    if isinstance(ds, int):
        if img.ndim == 3:
            res = tomopy.downsample(img, level=int(np.log2(ds)), axis=0)
        else:
            res = tomopy.downsample(img[:, np.newaxis, :], level=int(np.log2(ds)), axis=0)
    else:
        zm = np.ones(img.ndim)
        zm[axis] = 1. / ds
        res = zoom(img, zm)
    return res

def lateral_damp(img, length=50, sigma=None):

    ind = int(length / 2)
    mask = np.ones(img.shape)
    mask[:, :ind] = 0
    mask[:, -ind:] = 0
    if sigma is None:
        sigma = ind / 5.
    mask = gaussian_filter(mask, sigma=sigma)
    return img * mask

def pad_sinogram(sino, length):

    length = int(length)
    res = np.zeros([sino.shape[0], sino.shape[1] + length * 2])
    res[:, length:length+sino.shape[1]] = sino
    mean_left = np.mean(sino[:, :40], axis=1).reshape([sino.shape[0], 1])
    mean_right = np.mean(sino[:, -40:], axis=1).reshape([sino.shape[0], 1])
    res[:, :length] = mean_left
    res[:, -length:] = mean_right

    return res

def get_fov(n_scan, scanned_sino_width, mask_ratio_local):
    """
    Get field of view width that allows full coverage in local tomography.
    """
    fov = scanned_sino_width if n_scan == 1 else int(np.sqrt(2) * scanned_sino_width / (n_scan - 1 + np.sqrt(2))) + 10
    fov = int(fov / mask_ratio_local)
    if fov % 2 == 1:
        fov += 1
    return fov

def get_nscan_ps(f, gammaps, l):

    if f >= l:
        return 1
    else:
        return int(np.ceil((l - f) / (gammaps * f) + 1))

def get_nscan_os(f, fprime, l):

    if f >= l:
        return 1
    else:
        return int(np.ceil(np.sqrt(2) * l / fprime))


text:
import numpy as np
import utm
import cv2
from collections import defaultdict
import logging

logger = logging.getLogger('app.' + __name__)

class Fov:
    def __init__(self):
        logger.debug(f'Creating instance of Fov {self}')
        self.image_size = None
        self.horizontal_fov = None
        self.vertical_fov = None
        self.camera_matrix = None
        self.dist_coefficients = None
        self.camera_matrix = None
        self.dist_coefficients = None

    def set_image_size(self, width, height):
        self.image_size = (width, height)

    def set_camera_params(self, mat_file):
        logger.debug(f'Setting camera params')
        mat_contents = np.load(mat_file)
        self.camera_matrix = mat_contents['mtx']
        self.dist_coefficients = mat_contents['dist']
        self.horizontal_fov = mat_contents['fov_x'] * np.pi / 180
        self.vertical_fov = mat_contents['fov_y'] * np.pi / 180

    @staticmethod
    def roll(roll):
        return np.array([[np.cos(roll), 0, np.sin(roll)], [0, 1, 0], [-np.sin(roll), 0, np.cos(roll)]])

    @staticmethod
    def pitch(pitch):
        return np.array([[1, 0, 0], [0, np.cos(pitch), -np.sin(pitch)], [0, np.sin(pitch), np.cos(pitch)]])

    @staticmethod
    def yaw(yaw):
        return np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]])

    def rotation(self, yaw, pitch, roll):
        return np.matmul(self.yaw(yaw), np.matmul(self.pitch(pitch), self.roll(roll)))

    def get_unit_vector(self, image_point):
        if self.camera_matrix is not None:
            undist_point = cv2.undistortPoints(np.array([[image_point]], dtype=np.float32), self.camera_matrix, self.dist_coefficients, P=self.camera_matrix)[0][0]
        else:
            undist_point = image_point
        image_center = np.array([self.image_size[0]/2, self.image_size[1]/2])
        image_point_from_center = undist_point - image_center
        image_plane_width_in_meters = np.tan(self.horizontal_fov/2)*2
        image_plane_height_in_meters = np.tan(self.vertical_fov/2)*2
        x = image_point_from_center[0] / self.image_size[0] * image_plane_width_in_meters
        y = 1
        z = - image_point_from_center[1] / self.image_size[1] * image_plane_height_in_meters
        vector = np.array([x, y, z])
        return vector

    def get_horizon_and_world_corners(self, world_point_dict, yaw_pitch_roll):
        margin = 200
        image_plane_width_in_meters = np.tan(self.horizontal_fov / 2) * 2
        image_plane_height_in_meters = np.tan(self.vertical_fov / 2) * 2
        yaw_pitch_roll = (-yaw_pitch_roll[0], yaw_pitch_roll[1], yaw_pitch_roll[2])
        rotation_matrix = self.rotation(*yaw_pitch_roll)
        image_points = defaultdict(list)
        for direction, world_points in world_point_dict.items():
            temp_image_points_x = []
            temp_image_points_y = []
            for world_point in world_points:
                world_rotated_vector = np.matmul(np.transpose(rotation_matrix), world_point)
                if world_rotated_vector[1] >= 0:
                    vector = world_rotated_vector / world_rotated_vector[1]
                    image_point_x = int(vector[0] / image_plane_width_in_meters * self.image_size[0] + self.image_size[0]/2)
                    image_point_y = int(- vector[2] / image_plane_height_in_meters * self.image_size[1] + self.image_size[1]/2)
                    if -margin <= image_point_x <= self.image_size[0] + margin:
                        if -margin <= image_point_y <= self.image_size[1] + margin:
                            temp_image_points_x.append(image_point_x)
                            temp_image_points_y.append(image_point_y)
            left = min(temp_image_points_x, default=0)
            top = min(temp_image_points_y, default=0)
            for x, y in zip(temp_image_points_x, temp_image_points_y):
                image_points[direction].append({'x': x - left, 'y': y - top})
            image_points[direction + '_pos'].append({'top': top, 'left': left})
        return image_points

    def get_world_point(self, image_point, drone_height, yaw_pitch_roll, pos, return_zone=False):
        unit_vector = self.get_unit_vector(image_point)
        yaw_pitch_roll = (-yaw_pitch_roll[0], yaw_pitch_roll[1], yaw_pitch_roll[2])
        rotation_matrix = self.rotation(*yaw_pitch_roll)
        rotated_vector = np.matmul(rotation_matrix, unit_vector)
        ground_vector = rotated_vector / rotated_vector[2] * -drone_height
        east_north_zone = self.convert_gps(*pos)
        world_point = ground_vector[:2] + np.array(east_north_zone[:2])
        if return_zone:
            return world_point, east_north_zone[2:]
        else:
            return world_point

    def get_world_points(self, image_points, drone_height, yaw_pitch_roll, pos):
        world_points = []
        for image_point in image_points:
            world_point = self.get_world_point(image_point, drone_height, yaw_pitch_roll, pos)
            world_points.append(world_point)
        return world_points

    def get_gps_point(self, image_point, drone_height, yaw_pitch_roll, pos):
        world_point, zone = self.get_world_point(image_point, drone_height, yaw_pitch_roll, pos, True)
        lat, lon = self.convert_utm(world_point[0], world_point[1], zone)
        return lat, lon

    @staticmethod
    def convert_gps(lat, lon):
        east_north_zone = utm.from_latlon(lat, lon)
        return east_north_zone

    @staticmethod
    def convert_utm(east, north, zone):
        lat, lon = utm.to_latlon(east, north, *zone)
        return lat, lon

fov = Fov()

ext: py

max_stars_repo_name: henrikmidtiby/DroneVideoMeasure

max_issues_repo_name: henrikmidtiby/DroneVideoMeasure

max_forks_repo_name: henrikmidtiby/DroneVideoMeasure

reason: import numpy


text:
# Copyright (c) 2017-present, Facebook, Inc.
#
# 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.
##############################################################################
#
# Based on:
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Detectron config system.

This file specifies default config options for Detectron. You should not
change values in this file. Instead, you should write a config file (in yaml)
and use merge_cfg_from_file(yaml_file) to load it and override the default
options.

Most tools in the tools directory take a --cfg option to specify an override
file and an optional list of override (key, value) pairs:
 - See tools/{train,test}_net.py for example code that uses merge_cfg_from_file
 - See configs/*/*.yaml for example config files

Detectron supports a lot of different model types, each of which has a lot of
different options. The result is a HUGE set of configuration options.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from ast import literal_eval
from past.builtins import basestring
from utils.collections import AttrDict
import copy
import logging
import numpy as np
import os
import os.path as osp
import yaml

from utils.io import cache_url

logger = logging.getLogger(__name__)

__C = AttrDict()
# Consumers can get config by:
#   from core.config import cfg
cfg = __C

# Random note: avoid using '.ON' as a config key since yaml converts it to True;
# prefer 'ENABLED' instead

# ---------------------------------------------------------------------------- #
# Training options
# ---------------------------------------------------------------------------- #
__C.TRAIN = AttrDict()

# Initialize network with weights from this .pkl file
__C.TRAIN.WEIGHTS = b''

# Datasets to train on
# Available dataset list: datasets.dataset_catalog.DATASETS.keys()
# If multiple datasets are listed, the model is trained on their union
__C.TRAIN.DATASETS = ()

# Scales to use during training
# Each scale is the pixel size of an image's shortest side
# If multiple scales are listed, then one is selected uniformly at random for
# each training image (i.e., scale jitter data augmentation)
__C.TRAIN.SCALES = (600, )

# Max pixel size of the longest side of a scaled input image
__C.TRAIN.MAX_SIZE = 1000

# Images *per GPU* in the training minibatch
# Total images per minibatch = TRAIN.IMS_PER_BATCH * NUM_GPUS
__C.TRAIN.IMS_PER_BATCH = 2

# RoI minibatch size *per image* (number of regions of interest [ROIs])
# Total number of RoIs per training minibatch =
#   TRAIN.BATCH_SIZE_PER_IM * TRAIN.IMS_PER_BATCH * NUM_GPUS
# E.g., a common configuration is: 512 * 2 * 8 = 8192
__C.TRAIN.BATCH_SIZE_PER_IM = 64

# Target fraction of RoI minibatch that is labeled foreground (i.e. class > 0)
__C.TRAIN.FG_FRACTION = 0.25

# Overlap threshold for an RoI to be considered foreground (if >= FG_THRESH)
__C.TRAIN.FG_THRESH = 0.5

# Overlap threshold for an RoI to be considered background (class = 0 if
# overlap in [LO, HI))
__C.TRAIN.BG_THRESH_HI = 0.5
__C.TRAIN.BG_THRESH_LO = 0.0

# Use horizontally-flipped images during training?
__C.TRAIN.USE_FLIPPED = True

# Overlap required between an RoI and a ground-truth box in order for that
# (RoI, gt box) pair to be used as a bounding-box regression training example
__C.TRAIN.BBOX_THRESH = 0.5

# Snapshot (model checkpoint) period
# Divide by NUM_GPUS to determine actual period (e.g., 20000/8 => 2500 iters)
# to allow for linear training schedule scaling
__C.TRAIN.SNAPSHOT_ITERS = 20000

# Train using these proposals
# During training, all proposals specified in the file are used (no limit is
# applied)
# Proposal files must be in correspondence with the datasets listed in
# TRAIN.DATASETS
__C.TRAIN.PROPOSAL_FILES = ()

# Make minibatches from images that have similar aspect ratios (i.e. both
# tall and thin or both short and wide)
# This feature is critical for saving memory (and makes training slightly
# faster)
__C.TRAIN.ASPECT_GROUPING = True

# ---------------------------------------------------------------------------- #
# RPN training options
# ---------------------------------------------------------------------------- #

# Minimum overlap required between an anchor and ground-truth box for the
# (anchor, gt box) pair to be a positive example (IOU >= thresh ==> positive RPN
# example)
__C.TRAIN.RPN_POSITIVE_OVERLAP = 0.7

# Maximum overlap allowed between an anchor and ground-truth box for the
# (anchor, gt box) pair to be a negative examples (IOU < thresh ==> negative RPN
# example)
__C.TRAIN.RPN_NEGATIVE_OVERLAP = 0.3

# Target fraction of foreground (positive) examples per RPN minibatch
__C.TRAIN.RPN_FG_FRACTION = 0.5

# Total number of RPN examples per image
__C.TRAIN.RPN_BATCH_SIZE_PER_IM = 256

# NMS threshold used on RPN proposals (used during end-to-end training with RPN)
__C.TRAIN.RPN_NMS_THRESH = 0.7

# Number of top scoring RPN proposals to keep before applying NMS
# When FPN is used, this is *per FPN level* (not total)
__C.TRAIN.RPN_PRE_NMS_TOP_N = 12000

# Number of top scoring RPN proposals to keep after applying NMS
# This is the total number of RPN proposals produced (for both FPN and non-FPN
# cases)
__C.TRAIN.RPN_POST_NMS_TOP_N = 2000

# Remove RPN anchors that go outside the image by RPN_STRADDLE_THRESH pixels
# Set to -1 or a large value, e.g. 100000, to disable pruning anchors
__C.TRAIN.RPN_STRADDLE_THRESH = 0

# Proposal height and width both need to be greater than RPN_MIN_SIZE
# (at orig image scale; not scale used during training or inference)
__C.TRAIN.RPN_MIN_SIZE = 0

# Filter proposals that are inside of crowd regions by CROWD_FILTER_THRESH
# "Inside" is measured as: proposal-with-crowd intersection area divided by
# proposal area
__C.TRAIN.CROWD_FILTER_THRESH = 0.7

# Ignore ground-truth objects with area < this threshold
__C.TRAIN.GT_MIN_AREA = -1

# Freeze the backbone architecture during training if set to True
__C.TRAIN.FREEZE_CONV_BODY = False

# Training will resume from the latest snapshot (model checkpoint) found in the
# output directory
__C.TRAIN.AUTO_RESUME = True

################################################################################
# For seg-every-thing training
__C.TRAIN.TRAIN_MASK_HEAD_ONLY = False
__C.TRAIN.MRCNN_FILTER_LABELS = False
__C.TRAIN.MRCNN_LABELS_TO_KEEP = ()
################################################################################

# ---------------------------------------------------------------------------- #
# Data loader options
# ---------------------------------------------------------------------------- #
__C.DATA_LOADER = AttrDict()

# Number of Python threads to use for the data loader (warning: using too many
# threads can cause GIL-based interference with Python Ops leading to *slower*
# training; 4 seems to be the sweet spot in our experience)
__C.DATA_LOADER.NUM_THREADS = 4

# ---------------------------------------------------------------------------- #
# Inference ('test') options
# ---------------------------------------------------------------------------- #
__C.TEST = AttrDict()

# Initialize network with weights from this .pkl file
__C.TEST.WEIGHTS = b''

# Datasets to test on
# Available dataset list: datasets.dataset_catalog.DATASETS.keys()
# If multiple datasets are listed, testing is performed on each one sequentially
__C.TEST.DATASETS = ()

# Scale to use during testing
__C.TEST.SCALE = 600

# Max pixel size of the longest side of a scaled input image
__C.TEST.MAX_SIZE = 1000

# Overlap threshold used for non-maximum suppression (suppress boxes with
# IoU >= this threshold)
__C.TEST.NMS = 0.3

# Apply Fast R-CNN style bounding-box regression if True
__C.TEST.BBOX_REG = True

# Test using these proposal files (must correspond with TEST.DATASETS)
__C.TEST.PROPOSAL_FILES = ()

# Limit on the number of proposals per image used during inference
__C.TEST.PROPOSAL_LIMIT = 2000

# NMS threshold used on RPN proposals
__C.TEST.RPN_NMS_THRESH = 0.7

# Number of top scoring RPN proposals to keep before applying NMS
# When FPN is used, this is *per FPN level* (not total)
__C.TEST.RPN_PRE_NMS_TOP_N = 12000

# Number of top scoring RPN proposals to keep after applying NMS
# This is the total number of RPN proposals produced (for both FPN and non-FPN
# cases)
__C.TEST.RPN_POST_NMS_TOP_N = 2000

# Proposal height and width both need to be greater than RPN_MIN_SIZE
# (at orig image scale; not scale used during training or inference)
__C.TEST.RPN_MIN_SIZE = 0

# Maximum number of detections to return per image (100 is based on the limit
# established for the COCO dataset)
__C.TEST.DETECTIONS_PER_IM = 100

# Minimum score threshold (assuming scores in a [0, 1] range); a value chosen to
# balance obtaining high recall with not having too many low precision
# detections that will slow down inference post processing steps (like NMS)
__C.TEST.SCORE_THRESH = 0.05

# Save detection results files if True
# If false, results files are cleaned up (they can be large) after local
# evaluation
__C.TEST.COMPETITION_MODE = True

# Evaluate detections with the COCO json dataset eval code even if it's not the
# evaluation code for the dataset (e.g. evaluate PASCAL VOC results using the
# COCO API to get COCO style AP on PASCAL VOC)
__C.TEST.FORCE_JSON_DATASET_EVAL = False

# [Inferred value; do not set directly in a config]
# Indicates if precomputed proposals are used at test time
# Not set for 1-stage models and 2-stage models with RPN subnetwork enabled
__C.TEST.PRECOMPUTED_PROPOSALS = True

# ---------------------------------------------------------------------------- #
# Test-time augmentations for bounding box detection
# See configs/test_time_aug/e2e_mask_rcnn_R-50-FPN_2x.yaml for an example
# ---------------------------------------------------------------------------- #
__C.TEST.BBOX_AUG = AttrDict()

# Enable test-time augmentation for bounding box detection if True
__C.TEST.BBOX_AUG.ENABLED = False

# Heuristic used to combine predicted box scores
#   Valid options: ('ID', 'AVG', 'UNION')
__C.TEST.BBOX_AUG.SCORE_HEUR = b'UNION'

# Heuristic used to combine predicted box coordinates
#   Valid options: ('ID', 'AVG', 'UNION')
__C.TEST.BBOX_AUG.COORD_HEUR = b'UNION'

# Horizontal flip at the original scale (id transform)
__C.TEST.BBOX_AUG.H_FLIP = False

# Each scale is the pixel size of an image's shortest side
__C.TEST.BBOX_AUG.SCALES = ()

# Max pixel size of the longer side
__C.TEST.BBOX_AUG.MAX_SIZE = 4000

# Horizontal flip at each scale
__C.TEST.BBOX_AUG.SCALE_H_FLIP = False

# Apply scaling based on object size
__C.TEST.BBOX_AUG.SCALE_SIZE_DEP = False
__C.TEST.BBOX_AUG.AREA_TH_LO = 50**2
__C.TEST.BBOX_AUG.AREA_TH_HI = 180**2

# Each aspect ratio is relative to image width
__C.TEST.BBOX_AUG.ASPECT_RATIOS = ()

# Horizontal flip at each aspect ratio
__C.TEST.BBOX_AUG.ASPECT_RATIO_H_FLIP = False

# ---------------------------------------------------------------------------- #
# Test-time augmentations for mask detection
# See configs/test_time_aug/e2e_mask_rcnn_R-50-FPN_2x.yaml for an example
# ---------------------------------------------------------------------------- #
__C.TEST.MASK_AUG = AttrDict()

# Enable test-time augmentation for instance mask detection if True
__C.TEST.MASK_AUG.ENABLED = False

# Heuristic used to combine mask predictions
# SOFT prefix indicates that the computation is performed on soft masks
#   Valid options: ('SOFT_AVG', 'SOFT_MAX', 'LOGIT_AVG')
__C.TEST.MASK_AUG.HEUR = b'SOFT_AVG'

# Horizontal flip at the original scale (id transform)
__C.TEST.MASK_AUG.H_FLIP = False

# Each scale is the pixel size of an image's shortest side
__C.TEST.MASK_AUG.SCALES = ()

# Max pixel size of the longer side
__C.TEST.MASK_AUG.MAX_SIZE = 4000

# Horizontal flip at each scale
__C.TEST.MASK_AUG.SCALE_H_FLIP = False

# Apply scaling based on object size
__C.TEST.MASK_AUG.SCALE_SIZE_DEP = False
__C.TEST.MASK_AUG.AREA_TH = 180**2

# Each aspect ratio is relative to image width
__C.TEST.MASK_AUG.ASPECT_RATIOS = ()

# Horizontal flip at each aspect ratio
__C.TEST.MASK_AUG.ASPECT_RATIO_H_FLIP = False

# ---------------------------------------------------------------------------- #
# Test-augmentations for keypoints detection
# configs/test_time_aug/keypoint_rcnn_R-50-FPN_1x.yaml
# ---------------------------------------------------------------------------- #
__C.TEST.KPS_AUG = AttrDict()

# Enable test-time augmentation for keypoint detection if True
__C.TEST.KPS_AUG.ENABLED = False

# Heuristic used to combine keypoint predictions
#   Valid options: ('HM_AVG', 'HM_MAX')
__C.TEST.KPS_AUG.HEUR = b'HM_AVG'

# Horizontal flip at the original scale (id transform)
__C.TEST.KPS_AUG.H_FLIP = False

# Each scale is the pixel size of an image's shortest side
__C.TEST.KPS_AUG.SCALES = ()

# Max pixel size of the longer side
__C.TEST.KPS_AUG.MAX_SIZE = 4000

# Horizontal flip at each scale
__C.TEST.KPS_AUG.SCALE_H_FLIP = False

# Apply scaling based on object size
__C.TEST.KPS_AUG.SCALE_SIZE_DEP = False
__C.TEST.KPS_AUG.AREA_TH = 180**2

# Eeach aspect ratio is realtive to image width
__C.TEST.KPS_AUG.ASPECT_RATIOS = ()

# Horizontal flip at each aspect ratio
__C.TEST.KPS_AUG.ASPECT_RATIO_H_FLIP = False

# ---------------------------------------------------------------------------- #
# Soft NMS
# ---------------------------------------------------------------------------- #
__C.TEST.SOFT_NMS = AttrDict()

# Use soft NMS instead of standard NMS if set to True
__C.TEST.SOFT_NMS.ENABLED = False
# See soft NMS paper for definition of these options
__C.TEST.SOFT_NMS.METHOD = b'linear'
__C.TEST.SOFT_NMS.SIGMA = 0.5
# For the soft NMS overlap threshold, we simply use TEST.NMS

# ---------------------------------------------------------------------------- #
# Bounding box voting (from the Multi-Region CNN paper)
# ---------------------------------------------------------------------------- #
__C.TEST.BBOX_VOTE = AttrDict()

# Use box voting if set to True
__C.TEST.BBOX_VOTE.ENABLED = False

# We use TEST.NMS threshold for the NMS step. VOTE_TH overlap threshold
# is used to select voting boxes (IoU >= VOTE_TH) for each box that survives NMS
__C.TEST.BBOX_VOTE.VOTE_TH = 0.8

# The method used to combine scores when doing bounding box voting
# Valid options include ('ID', 'AVG', 'IOU_AVG', 'GENERALIZED_AVG', 'QUASI_SUM')
__C.TEST.BBOX_VOTE.SCORING_METHOD = b'ID'

# Hyperparameter used by the scoring method (it has different meanings for
# different methods)
__C.TEST.BBOX_VOTE.SCORING_METHOD_BETA = 1.0

# ---------------------------------------------------------------------------- #
# Model options
# ---------------------------------------------------------------------------- #
__C.MODEL = AttrDict()

# The type of model to use
# The string must match a function in the modeling.model_builder module
# (e.g., 'generalized_rcnn', 'mask_rcnn', ...)
__C.MODEL.TYPE = b''

# The backbone conv body to use
# The string must match a function that is imported in modeling.model_builder
# (e.g., 'FPN.add_fpn_ResNet101_conv5_body' to specify a ResNet-101-FPN
# backbone)
__C.MODEL.CONV_BODY = b''

# Number of classes in the dataset; must be set
# E.g., 81 for COCO (80 foreground + 1 background)
__C.MODEL.NUM_CLASSES = -1

# Use a class agnostic bounding box regressor instead of the default per-class
# regressor
__C.MODEL.CLS_AGNOSTIC_BBOX_REG = False

# Default weights on (dx, dy, dw, dh) for normalizing bbox regression targets
# These are empirically chosen to approximately lead to unit variance targets
__C.MODEL.BBOX_REG_WEIGHTS = (10., 10., 5., 5.)

# The meaning of FASTER_RCNN depends on the context (training vs. inference):
# 1) During training, FASTER_RCNN = True means that end-to-end training will be
#    used to jointly train the RPN subnetwork and the Fast R-CNN subnetwork
#    (Faster R-CNN = RPN + Fast R-CNN).
# 2) During inference, FASTER_RCNN = True means that the model's RPN subnetwork
#    will be used to generate proposals rather than relying on precomputed
#    proposals. Note that FASTER_RCNN = True can be used at inference time even
#    if the Faster R-CNN model was trained with stagewise training (which
#    consists of alternating between RPN and Fast R-CNN training in a way that
#    finally leads to a single network).
__C.MODEL.FASTER_RCNN = False

# Indicates the model makes instance mask predictions (as in Mask R-CNN)
__C.MODEL.MASK_ON = False

# Indicates the model makes keypoint predictions (as in Mask R-CNN for
# keypoints)
__C.MODEL.KEYPOINTS_ON = False

# Indicates the model's computation terminates with the production of RPN
# proposals (i.e., it outputs proposals ONLY, no actual object detections)
__C.MODEL.RPN_ONLY = False

# Caffe2 net execution type
# Use 'prof_dag' to get profiling statistics
__C.MODEL.EXECUTION_TYPE = b'dag'

# ---------------------------------------------------------------------------- #
# RetinaNet options
# ---------------------------------------------------------------------------- #
__C.RETINANET = AttrDict()

# RetinaNet is used (instead of Fast/er/Mask R-CNN/R-FCN/RPN) if True
__C.RETINANET.RETINANET_ON = False

# Anchor aspect ratios to use
__C.RETINANET.ASPECT_RATIOS = (0.5, 1.0, 2.0)

# Anchor scales per octave
__C.RETINANET.SCALES_PER_OCTAVE = 3

# At each FPN level, we generate anchors based on their scale, aspect_ratio,
# stride of the level, and we multiply the resulting anchor by ANCHOR_SCALE
__C.RETINANET.ANCHOR_SCALE = 4

# Convolutions to use in the cls and bbox tower
# NOTE: this doesn't include the last conv for logits
__C.RETINANET.NUM_CONVS = 4

# Weight for bbox_regression loss
__C.RETINANET.BBOX_REG_WEIGHT = 1.0

# Smooth L1 loss beta for bbox regression
__C.RETINANET.BBOX_REG_BETA = 0.11

# During inference, #locs to select based on cls score before NMS is performed
# per FPN level
__C.RETINANET.PRE_NMS_TOP_N = 1000

# IoU overlap ratio for labeling an anchor as positive
# Anchors with >= iou overlap are labeled positive
__C.RETINANET.POSITIVE_OVERLAP = 0.5

# IoU overlap ratio for labeling an anchor as negative
# Anchors with < iou overlap are labeled negative
__C.RETINANET.NEGATIVE_OVERLAP = 0.4

# Focal loss parameter: alpha
__C.RETINANET.LOSS_ALPHA = 0.25

# Focal loss parameter: gamma
__C.RETINANET.LOSS_GAMMA = 2.0

# Prior prob for the positives at the beginning of training. This is used to set
# the bias init for the logits layer
__C.RETINANET.PRIOR_PROB = 0.01

# Whether classification and bbox branch tower should be shared or not
__C.RETINANET.SHARE_CLS_BBOX_TOWER = False

# Use class specific bounding box regression instead of the default class
# agnostic regression
__C.RETINANET.CLASS_SPECIFIC_BBOX = False

# Whether softmax should be used in classification branch training
__C.RETINANET.SOFTMAX = False

# Inference cls score threshold, anchors with score > INFERENCE_TH are
# considered for inference
__C.RETINANET.INFERENCE_TH = 0.05

# ---------------------------------------------------------------------------- #
# Solver options
# Note: all solver options are used exactly as specified; the implication is
# that if you switch from training on 1 GPU to N GPUs, you MUST adjust the
# solver configuration accordingly. We suggest using gradual warmup and the
# linear learning rate scaling rule as described in
# "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" Goyal et al.
# https://arxiv.org/abs/1706.02677
# ---------------------------------------------------------------------------- #
__C.SOLVER = AttrDict()

# Base learning rate for the specified schedule
__C.SOLVER.BASE_LR = 0.001

# Schedule type (see functions in utils.lr_policy for options)
# E.g., 'step', 'steps_with_decay', ...
__C.SOLVER.LR_POLICY = b'step'

# Some LR Policies (by example):
# 'step'
#   lr = SOLVER.BASE_LR * SOLVER.GAMMA ** (cur_iter // SOLVER.STEP_SIZE)
# 'steps_with_decay'
#   SOLVER.STEPS = [0, 60000, 80000]
#   SOLVER.GAMMA = 0.1
#   lr = SOLVER.BASE_LR * SOLVER.GAMMA ** current_step
#   iters [0, 59999] are in current_step = 0, iters [60000, 79999] are in
#   current_step = 1, and so on
# 'steps_with_lrs'
#   SOLVER.STEPS = [0, 60000, 80000]
#   SOLVER.LRS = [0.02, 0.002, 0.0002]
#   lr = LRS[current_step]

# Hyperparameter used by the specified policy
# For 'step', the current LR is multiplied by SOLVER.GAMMA at each step
__C.SOLVER.GAMMA = 0.1

# Uniform step size for 'steps' policy
__C.SOLVER.STEP_SIZE = 30000

# Non-uniform step iterations for 'steps_with_decay' or 'steps_with_lrs'
# policies
__C.SOLVER.STEPS = []

# Learning rates to use with 'steps_with_lrs' policy
__C.SOLVER.LRS = []

# Maximum number of SGD iterations
__C.SOLVER.MAX_ITER = 40000

# Momentum to use with SGD
__C.SOLVER.MOMENTUM = 0.9

# L2 regularization hyperparameter
__C.SOLVER.WEIGHT_DECAY = 0.0005
# L2 regularization hyperparameter for GroupNorm's parameters
__C.SOLVER.WEIGHT_DECAY_GN = 0.0

# Warm up to SOLVER.BASE_LR over this number of SGD iterations
__C.SOLVER.WARM_UP_ITERS = 500

# Start the warm up from SOLVER.BASE_LR * SOLVER.WARM_UP_FACTOR
__C.SOLVER.WARM_UP_FACTOR = 1.0 / 3.0

# WARM_UP_METHOD can be either 'constant' or 'linear' (i.e., gradual)
__C.SOLVER.WARM_UP_METHOD = 'linear'

# Scale the momentum update history by new_lr / old_lr when updating the
# learning rate (this is correct given MomentumSGDUpdateOp)
__C.SOLVER.SCALE_MOMENTUM = True
# Only apply the correction if the relative LR change exceeds this threshold
# (prevents ever change in linear warm up from scaling the momentum by a tiny
# amount; momentum scaling is only important if the LR change is large)
__C.SOLVER.SCALE_MOMENTUM_THRESHOLD = 1.1

# Suppress logging of changes to LR unless the relative change exceeds this
# threshold (prevents linear warm up from spamming the training log)
__C.SOLVER.LOG_LR_CHANGE_THRESHOLD = 1.1

# ---------------------------------------------------------------------------- #
# Fast R-CNN options
# ---------------------------------------------------------------------------- #
__C.FAST_RCNN = AttrDict()

# The type of RoI head to use for bounding box classification and regression
# The string must match a function this is imported in modeling.model_builder
# (e.g., 'head_builder.add_roi_2mlp_head' to specify a two hidden layer MLP)
__C.FAST_RCNN.ROI_BOX_HEAD = b''

# Hidden layer dimension when using an MLP for the RoI box head
__C.FAST_RCNN.MLP_HEAD_DIM = 1024

# Hidden Conv layer dimension when using Convs for the RoI box head
__C.FAST_RCNN.CONV_HEAD_DIM = 256
# Number of stacked Conv layers in the RoI box head
__C.FAST_RCNN.NUM_STACKED_CONVS = 4

# RoI transformation function (e.g., RoIPool or RoIAlign)
# (RoIPoolF is the same as RoIPool; ignore the trailing 'F')
__C.FAST_RCNN.ROI_XFORM_METHOD = b'RoIPoolF'

# Number of grid sampling points in RoIAlign (usually use 2)
# Only applies to RoIAlign
__C.FAST_RCNN.ROI_XFORM_SAMPLING_RATIO = 0

# RoI transform output resolution
# Note: some models may have constraints on what they can use, e.g. they use
# pretrained FC layers like in VGG16, and will ignore this option
__C.FAST_RCNN.ROI_XFORM_RESOLUTION = 14

# ---------------------------------------------------------------------------- #
# RPN options
# ---------------------------------------------------------------------------- #
__C.RPN = AttrDict()

# [Infered value; do not set directly in a config]
# Indicates that the model contains an RPN subnetwork
__C.RPN.RPN_ON = False

# RPN anchor sizes given in absolute pixels w.r.t. the scaled network input
# Note: these options are *not* used by FPN RPN; see FPN.RPN* options
__C.RPN.SIZES = (64, 128, 256, 512)

# Stride of the feature map that RPN is attached
__C.RPN.STRIDE = 16

# RPN anchor aspect ratios
__C.RPN.ASPECT_RATIOS = (0.5, 1, 2)

# ---------------------------------------------------------------------------- #
# FPN options
# ---------------------------------------------------------------------------- #
__C.FPN = AttrDict()

# FPN is enabled if True
__C.FPN.FPN_ON = False

# Channel dimension of the FPN feature levels
__C.FPN.DIM = 256

# Initialize the lateral connections to output zero if True
__C.FPN.ZERO_INIT_LATERAL = False

# Stride of the coarsest FPN level
# This is needed so the input can be padded properly
__C.FPN.COARSEST_STRIDE = 32

#
# FPN may be used for just RPN, just object detection, or both
#

# Use FPN for RoI transform for object detection if True
__C.FPN.MULTILEVEL_ROIS = False
# Hyperparameters for the RoI-to-FPN level mapping heuristic
__C.FPN.ROI_CANONICAL_SCALE = 224  # s0
__C.FPN.ROI_CANONICAL_LEVEL = 4  # k0: where s0 maps to
# Coarsest level of the FPN pyramid
__C.FPN.ROI_MAX_LEVEL = 5
# Finest level of the FPN pyramid
__C.FPN.ROI_MIN_LEVEL = 2

# Use FPN for RPN if True
__C.FPN.MULTILEVEL_RPN = False
# Coarsest level of the FPN pyramid
__C.FPN.RPN_MAX_LEVEL = 6
# Finest level of the FPN pyramid
__C.FPN.RPN_MIN_LEVEL = 2
# FPN RPN anchor aspect ratios
__C.FPN.RPN_ASPECT_RATIOS = (0.5, 1, 2)
# RPN anchors start at this size on RPN_MIN_LEVEL
# The anchor size doubled each level after that
# With a default of 32 and levels 2 to 6, we get anchor sizes of 32 to 512
__C.FPN.RPN_ANCHOR_START_SIZE = 32
# Use extra FPN levels, as done in the RetinaNet paper
__C.FPN.EXTRA_CONV_LEVELS = False
# Use GroupNorm in the FPN-specific layers (lateral, etc.)
__C.FPN.USE_GN = False

# ---------------------------------------------------------------------------- #
# Mask R-CNN options ("MRCNN" means Mask R-CNN)
# ---------------------------------------------------------------------------- #
__C.MRCNN = AttrDict()

# The type of RoI head to use for instance mask prediction
# The string must match a function this is imported in modeling.model_builder
# (e.g., 'mask_rcnn_heads.ResNet_mask_rcnn_fcn_head_v1up4convs')
__C.MRCNN.ROI_MASK_HEAD = b''

# Resolution of mask predictions
__C.MRCNN.RESOLUTION = 14

# RoI transformation function and associated options
__C.MRCNN.ROI_XFORM_METHOD = b'RoIAlign'

# RoI transformation function (e.g., RoIPool or RoIAlign)
__C.MRCNN.ROI_XFORM_RESOLUTION = 7

# Number of grid sampling points in RoIAlign (usually use 2)
# Only applies to RoIAlign
__C.MRCNN.ROI_XFORM_SAMPLING_RATIO = 0

# Number of channels in the mask head
__C.MRCNN.DIM_REDUCED = 256

# Use dilated convolution in the mask head
__C.MRCNN.DILATION = 2

# Upsample the predicted masks by this factor
__C.MRCNN.UPSAMPLE_RATIO = 1

# Use a fully-connected layer to predict the final masks instead of a conv layer
__C.MRCNN.USE_FC_OUTPUT = False

# Weight initialization method for the mask head and mask output layers
__C.MRCNN.CONV_INIT = b'GaussianFill'

# Use class specific mask predictions if True (otherwise use class agnostic mask
# predictions)
__C.MRCNN.CLS_SPECIFIC_MASK = True

# Multi-task loss weight for masks
__C.MRCNN.WEIGHT_LOSS_MASK = 1.0

# Binarization threshold for converting soft masks to hard masks
__C.MRCNN.THRESH_BINARIZE = 0.5

# ---------------------------------------------------------------------------- #
# Keyoint Mask R-CNN options ("KRCNN" = Mask R-CNN with Keypoint support)
# ---------------------------------------------------------------------------- #
__C.KRCNN = AttrDict()

# The type of RoI head to use for instance keypoint prediction
# The string must match a function this is imported in modeling.model_builder
# (e.g., 'keypoint_rcnn_heads.add_roi_pose_head_v1convX')
__C.KRCNN.ROI_KEYPOINTS_HEAD = b''

# Output size (and size loss is computed on), e.g., 56x56
__C.KRCNN.HEATMAP_SIZE = -1

# Use bilinear interpolation to upsample the final heatmap by this factor
__C.KRCNN.UP_SCALE = -1

# Apply a ConvTranspose layer to the hidden representation computed by the
# keypoint head prior to predicting the per-keypoint heatmaps
__C.KRCNN.USE_DECONV = False
# Channel dimension of the hidden representation produced by the ConvTranspose
__C.KRCNN.DECONV_DIM = 256

# Use a ConvTranspose layer to predict the per-keypoint heatmaps
__C.KRCNN.USE_DECONV_OUTPUT = False

# Use dilation in the keypoint head
__C.KRCNN.DILATION = 1

# Size of the kernels to use in all ConvTranspose operations
__C.KRCNN.DECONV_KERNEL = 4

# Number of keypoints in the dataset (e.g., 17 for COCO)
__C.KRCNN.NUM_KEYPOINTS = -1

# Number of stacked Conv layers in keypoint head
__C.KRCNN.NUM_STACKED_CONVS = 8

# Dimension of the hidden representation output by the keypoint head
__C.KRCNN.CONV_HEAD_DIM = 256

# Conv kernel size used in the keypoint head
__C.KRCNN.CONV_HEAD_KERNEL = 3
# Conv kernel weight filling function
__C.KRCNN.CONV_INIT = b'GaussianFill'

# Use NMS based on OKS if True
__C.KRCNN.NMS_OKS = False

# Source of keypoint confidence
#   Valid options: ('bbox', 'logit', 'prob')
__C.KRCNN.KEYPOINT_CONFIDENCE = b'bbox'

# Standard ROI XFORM options (see FAST_RCNN or MRCNN options)
__C.KRCNN.ROI_XFORM_METHOD = b'RoIAlign'
__C.KRCNN.ROI_XFORM_RESOLUTION = 7
__C.KRCNN.ROI_XFORM_SAMPLING_RATIO = 0

# Minimum number of labeled keypoints that must exist in a minibatch (otherwise
# the minibatch is discarded)
__C.KRCNN.MIN_KEYPOINT_COUNT_FOR_VALID_MINIBATCH = 20

# When infering the keypoint locations from the heatmap, don't scale the heatmap
# below this minimum size
__C.KRCNN.INFERENCE_MIN_SIZE = 0

# Multi-task loss weight to use for keypoints
# Recommended values:
#   - use 1.0 if KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS is True
#   - use 4.0 if KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS is False
__C.KRCNN.LOSS_WEIGHT = 1.0

# Normalize by the total number of visible keypoints in the minibatch if True.
# Otherwise, normalize by the total number of keypoints that could ever exist
# in the minibatch. See comments in modeling.model_builder.add_keypoint_losses
# for detailed discussion.
__C.KRCNN.NORMALIZE_BY_VISIBLE_KEYPOINTS = True

################################################################################
# For seg-every-thing training
__C.MRCNN.BBOX2MASK = AttrDict()
__C.MRCNN.BBOX2MASK.BBOX2MASK_ON = False
__C.MRCNN.BBOX2MASK.TYPE = b''
__C.MRCNN.BBOX2MASK.USE_PRETRAINED_EMBED = False
__C.MRCNN.BBOX2MASK.PRETRAINED_EMBED_NAME = b''
__C.MRCNN.BBOX2MASK.PRETRAINED_EMBED_DIM = -1
__C.MRCNN.BBOX2MASK.STOP_DET_W_GRAD = True
__C.MRCNN.BBOX2MASK.INCLUDE_CLS_SCORE = True
__C.MRCNN.BBOX2MASK.INCLUDE_BBOX_PRED = False
__C.MRCNN.BBOX2MASK.USE_LEAKYRELU = True

__C.MRCNN.JOINT_FCN_MLP_HEAD = False
__C.MRCNN.MLP_MASK_BRANCH_TYPE = b''
################################################################################

# ---------------------------------------------------------------------------- #
# R-FCN options
# ---------------------------------------------------------------------------- #
__C.RFCN = AttrDict()

# Position-sensitive RoI pooling output grid size (height and width)
__C.RFCN.PS_GRID_SIZE = 3

# ---------------------------------------------------------------------------- #
# ResNets options ("ResNets" = ResNet and ResNeXt)
# ---------------------------------------------------------------------------- #
__C.RESNETS = AttrDict()

# Number of groups to use; 1 ==> ResNet; > 1 ==> ResNeXt
__C.RESNETS.NUM_GROUPS = 1

# Baseline width of each group
__C.RESNETS.WIDTH_PER_GROUP = 64

# Place the stride 2 conv on the 1x1 filter
# Use True only for the original MSRA ResNet; use False for C2 and Torch models
__C.RESNETS.STRIDE_1X1 = True

# Residual transformation function
__C.RESNETS.TRANS_FUNC = b'bottleneck_transformation'
# ResNet's stem function (conv1 and pool1)
__C.RESNETS.STEM_FUNC = b'basic_bn_stem'
# ResNet's shortcut function
__C.RESNETS.SHORTCUT_FUNC = b'basic_bn_shortcut'

# Apply dilation in stage "res5"
__C.RESNETS.RES5_DILATION = 1

# ---------------------------------------------------------------------------- #
# GroupNorm options
# ---------------------------------------------------------------------------- #
__C.GROUP_NORM = AttrDict()
# Number of dimensions per group in GroupNorm (-1 if using NUM_GROUPS)
__C.GROUP_NORM.DIM_PER_GP = -1
# Number of groups in GroupNorm (-1 if using DIM_PER_GP)
__C.GROUP_NORM.NUM_GROUPS = 32
# GroupNorm's small constant in the denominator
__C.GROUP_NORM.EPSILON = 1e-5

# ---------------------------------------------------------------------------- #
# Misc options
# ---------------------------------------------------------------------------- #

# Number of GPUs to use (applies to both training and testing)
__C.NUM_GPUS = 1

# Use NCCL for all reduce, otherwise use muji
# Warning: if set to True, you may experience deadlocks
__C.USE_NCCL = False

# The mapping from image coordinates to feature map coordinates might cause
# some boxes that are distinct in image space to become identical in feature
# coordinates. If DEDUP_BOXES > 0, then DEDUP_BOXES is used as the scale factor
# for identifying duplicate boxes.
# 1/16 is correct for {Alex,Caffe}Net, VGG_CNN_M_1024, and VGG16
__C.DEDUP_BOXES = 1 / 16.

# Clip bounding box transformation predictions to prevent np.exp from
# overflowing
# Heuristic choice based on that would scale a 16 pixel anchor up to 1000 pixels
__C.BBOX_XFORM_CLIP = np.log(1000. / 16.)

# Pixel mean values (BGR order) as a (1, 1, 3) array
# We use the same pixel mean for all networks even though it's not exactly what
# they were trained with
# "Fun" fact: the history of where these values comes from is lost
__C.PIXEL_MEANS = np.array([[[102.9801, 115.9465, 122.7717]]])

# For reproducibility...but not really because modern fast GPU libraries use
# non-deterministic op implementations
__C.RNG_SEED = 3

# A small number that's used many times
__C.EPS = 1e-14

# Root directory of project
__C.ROOT_DIR = os.getcwd()

# Output basedir
__C.OUTPUT_DIR = b'/tmp'

# Name (or path to) the matlab executable
__C.MATLAB = b'matlab'

# Reduce memory usage with memonger gradient blob sharing
__C.MEMONGER = True

# Futher reduce memory by allowing forward pass activations to be shared when
# possible. Note that this will cause activation blob inspection (values,
# shapes, etc.) to be meaningless when activation blobs are reused.
__C.MEMONGER_SHARE_ACTIVATIONS = False

# Dump detection visualizations
__C.VIS = False

# Score threshold for visualization
__C.VIS_TH = 0.9

# Expected results should take the form of a list of expectations, each
# specified by four elements (dataset, task, metric, expected value). For
# example: [['coco_2014_minival', 'box_proposal', 'AR@1000', 0.387]]
__C.EXPECTED_RESULTS = []
# Absolute and relative tolerance to use when comparing to EXPECTED_RESULTS
__C.EXPECTED_RESULTS_RTOL = 0.1
__C.EXPECTED_RESULTS_ATOL = 0.005
# Set to send email in case of an EXPECTED_RESULTS failure
__C.EXPECTED_RESULTS_EMAIL = b''

# Models and proposals referred to by URL are downloaded to a local cache
# specified by DOWNLOAD_CACHE
__C.DOWNLOAD_CACHE = b'/tmp/detectron-download-cache'

# ---------------------------------------------------------------------------- #
# Cluster options
# ---------------------------------------------------------------------------- #
__C.CLUSTER = AttrDict()

# Flag to indicate if the code is running in a cluster environment
__C.CLUSTER.ON_CLUSTER = False

# ---------------------------------------------------------------------------- #
# Deprecated options
# If an option is removed from the code and you don't want to break existing
# yaml configs, you can add the full config key as a string to the set below.
# ---------------------------------------------------------------------------- #
_DEPCRECATED_KEYS = set(
    {
        'FINAL_MSG',
        'MODEL.DILATION',
        'ROOT_GPU_ID',
        'RPN.ON',
        'TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED',
        'TRAIN.DROPOUT',
        'USE_GPU_NMS',
        'TEST.NUM_TEST_IMAGES',
    }
)

# ---------------------------------------------------------------------------- #
# Renamed options
# If you rename a config option, record the mapping from the old name to the new
# name in the dictionary below. Optionally, if the type also changed, you can
# make the value a tuple that specifies first the renamed key and then
# instructions for how to edit the config file.
# ---------------------------------------------------------------------------- #
_RENAMED_KEYS = {
    'EXAMPLE.RENAMED.KEY': 'EXAMPLE.KEY',  # Dummy example to follow
    'MODEL.PS_GRID_SIZE': 'RFCN.PS_GRID_SIZE',
    'MODEL.ROI_HEAD': 'FAST_RCNN.ROI_BOX_HEAD',
    'MRCNN.MASK_HEAD_NAME': 'MRCNN.ROI_MASK_HEAD',
    'TRAIN.DATASET': (
        'TRAIN.DATASETS',
        "Also convert to a tuple, e.g., " +
        "'coco_2014_train' -> ('coco_2014_train',) or " +
        "'coco_2014_train:coco_2014_valminusminival' -> " +
        "('coco_2014_train', 'coco_2014_valminusminival')"
    ),
    'TRAIN.PROPOSAL_FILE': (
        'TRAIN.PROPOSAL_FILES',
        "Also convert to a tuple, e.g., " +
        "'path/to/file' -> ('path/to/file',) or " +
        "'path/to/file1:path/to/file2' -> " +
        "('path/to/file1', 'path/to/file2')"
    ),
    'TEST.SCALES': (
        'TEST.SCALE',
        "Also convert from a tuple, e.g. (600, ), " +
        "to a integer, e.g. 600."
    ),
    'TEST.DATASET': (
        'TEST.DATASETS',
        "Also convert from a string, e.g 'coco_2014_minival', " +
        "to a tuple, e.g. ('coco_2014_minival', )."
    ),
    'TEST.PROPOSAL_FILE': (
        'TEST.PROPOSAL_FILES',
        "Also convert from a string, e.g. '/path/to/props.pkl', " +
        "to a tuple, e.g. ('/path/to/props.pkl', )."
    ),
}

def assert_and_infer_cfg(cache_urls=True, make_immutable=True):
    """Call this function in your script after you have finished setting all cfg
    values that are necessary (e.g., merging a config from a file, merging
    command line config options, etc.). By default, this function will also
    mark the global cfg as immutable to prevent changing the global cfg settings
    during script execution (which can lead to hard to debug errors or code
    that's harder to understand than is necessary).
    """
    if __C.MODEL.RPN_ONLY or __C.MODEL.FASTER_RCNN:
        __C.RPN.RPN_ON = True
    if __C.RPN.RPN_ON or __C.RETINANET.RETINANET_ON:
        __C.TEST.PRECOMPUTED_PROPOSALS = False
    if cache_urls:
        cache_cfg_urls()
    if make_immutable:
        cfg.immutable(True)

def cache_cfg_urls():
    """Download URLs in the config, cache them locally, and rewrite cfg to make
    use of the locally cached file.
    """
    __C.TRAIN.WEIGHTS = cache_url(__C.TRAIN.WEIGHTS, __C.DOWNLOAD_CACHE)
    __C.TEST.WEIGHTS = cache_url(__C.TEST.WEIGHTS, __C.DOWNLOAD_CACHE)
    __C.TRAIN.PROPOSAL_FILES = tuple(
        cache_url(f, __C.DOWNLOAD_CACHE) for f in __C.TRAIN.PROPOSAL_FILES
    )
    __C.TEST.PROPOSAL_FILES = tuple(
        cache_url(f, __C.DOWNLOAD_CACHE) for f in __C.TEST.PROPOSAL_FILES
    )

def get_output_dir(datasets, training=True):
    """Get the output directory determined by the current global config."""
    assert isinstance(datasets, (tuple, list, basestring)), \
        'datasets argument must be of type tuple, list or string'
    is_string = isinstance(datasets, basestring)
    dataset_name = datasets if is_string else ':'.join(datasets)
    tag = 'train' if training else 'test'
    # <output-dir>/<train|test>/<dataset-name>/<model-type>/
    outdir = osp.join(__C.OUTPUT_DIR, tag, dataset_name, __C.MODEL.TYPE)
    if not osp.exists(outdir):
        os.makedirs(outdir)
    return outdir

def merge_cfg_from_file(cfg_filename):
    """Load a yaml config file and merge it into the global config."""
    with open(cfg_filename, 'r') as f:
        yaml_cfg = AttrDict(yaml.load(f))
    _merge_a_into_b(yaml_cfg, __C)

def merge_cfg_from_cfg(cfg_other):
    """Merge `cfg_other` into the global config."""
    _merge_a_into_b(cfg_other, __C)

def merge_cfg_from_list(cfg_list):
    """Merge config keys, values in a list (e.g., from command line) into the
    global config. For example, `cfg_list = ['TEST.NMS', 0.5]`.
    """
    assert len(cfg_list) % 2 == 0
    for full_key, v in zip(cfg_list[0::2], cfg_list[1::2]):
        if _key_is_deprecated(full_key):
            continue
        if _key_is_renamed(full_key):
            _raise_key_rename_error(full_key)
        key_list = full_key.split('.')
        d = __C
        for subkey in key_list[:-1]:
            assert subkey in d, 'Non-existent key: {}'.format(full_key)
            d = d[subkey]
        subkey = key_list[-1]
        assert subkey in d, 'Non-existent key: {}'.format(full_key)
        value = _decode_cfg_value(v)
        value = _check_and_coerce_cfg_value_type(
            value, d[subkey], subkey, full_key
        )
        d[subkey] = value

def _merge_a_into_b(a, b, stack=None):
    """Merge config dictionary a into config dictionary b, clobbering the
    options in b whenever they are also specified in a.
    """
    assert isinstance(a, AttrDict), 'Argument `a` must be an AttrDict'
    assert isinstance(b, AttrDict), 'Argument `b` must be an AttrDict'

    for k, v_ in a.items():
        full_key = '.'.join(stack) + '.' + k if stack is not None else k
        # a must specify keys that are in b
        if k not in b:
            if _key_is_deprecated(full_key):
                continue
            elif _key_is_renamed(full_key):
                _raise_key_rename_error(full_key)
            else:
                raise KeyError('Non-existent config key: {}'.format(full_key))

        v = copy.deepcopy(v_)
        v = _decode_cfg_value(v)
        v = _check_and_coerce_cfg_value_type(v, b[k], k, full_key)

        # Recursively merge dicts
        if isinstance(v, AttrDict):
            try:
                stack_push = [k] if stack is None else stack + [k]
                _merge_a_into_b(v, b[k], stack=stack_push)
            except BaseException:
                raise
        else:
            b[k] = v

def _key_is_deprecated(full_key):
    if full_key in _DEPCRECATED_KEYS:
        logger.warn(
            'Deprecated config key (ignoring): {}'.format(full_key)
        )
        return True
    return False

def _key_is_renamed(full_key):
    return full_key in _RENAMED_KEYS

def _raise_key_rename_error(full_key):
    new_key = _RENAMED_KEYS[full_key]
    if isinstance(new_key, tuple):
        msg = ' Note: ' + new_key[1]
        new_key = new_key[0]
    else:
        msg = ''
    raise KeyError(
        'Key {} was renamed to {}; please update your config.{}'.
        format(full_key, new_key, msg)
    )

def _decode_cfg_value(v):
    """Decodes a raw config value (e.g., from a yaml config files or command
    line argument) into a Python object.
    """
    # Configs parsed from raw yaml will contain dictionary keys that need to be
    # converted to AttrDict objects
    if isinstance(v, dict):
        return AttrDict(v)
    # All remaining processing is only applied to strings
    if not isinstance(v, basestring):
        return v
    # Try to interpret `v` as a:
    #   string, number, tuple, list, dict, boolean, or None
    try:
        v = literal_eval(v)
    # The following two excepts allow v to pass through when it represents a
    # string.
    #
    # Longer explanation:
    # The type of v is always a string (before calling literal_eval), but
    # sometimes it *represents* a string and other times a data structure, like
    # a list. In the case that v represents a string, what we got back from the
    # yaml parser is 'foo' *without quotes* (so, not '"foo"'). literal_eval is
    # ok with '"foo"', but will raise a ValueError if given 'foo'. In other
    # cases, like paths (v = 'foo/bar' and not v = '"foo/bar"'), literal_eval
    # will raise a SyntaxError.
    except ValueError:
        pass
    except SyntaxError:
        pass
    return v

def _check_and_coerce_cfg_value_type(value_a, value_b, key, full_key):
    """Checks that `value_a`, which is intended to replace `value_b` is of the
    right type. The type is correct if it matches exactly or is one of a few
    cases in which the type can be easily coerced.
    """
    # The types must match (with some exceptions)
    type_b = type(value_b)
    type_a = type(value_a)
    if type_a is type_b:
        return value_a

    # Exceptions: numpy arrays, strings, tuple<->list
    if isinstance(value_b, np.ndarray):
        value_a = np.array(value_a, dtype=value_b.dtype)
    elif isinstance(value_b, basestring):
        value_a = str(value_a)
    elif isinstance(value_a, tuple) and isinstance(value_b, list):
        value_a = list(value_a)
    elif isinstance(value_a, list) and isinstance(value_b, tuple):
        value_a = tuple(value_a)
    else:
        raise ValueError(
            'Type mismatch ({} vs. {}) with values ({} vs. {}) for config '
            'key: {}'.format(type_b, type_a, value_b, value_a, full_key)
        )
    return value_a


text:
import numpy as np
from pyrr import aabb, line, ray
from pqdict import pqdict

class MyPlanner:
    __slots__ = ['boundary', 'blocks', 'res']

    def __init__(self, boundary, blocks, res):
        self.boundary = boundary
        self.blocks = blocks
        self.res = res

    def plan(self, start, goal):
        path = [start]
        numofdirs = 26
        [dX, dY, dZ] = np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])
        dR = np.vstack((dX.flatten(), dY.flatten(), dZ.flatten()))
        dR = np.delete(dR, 13, axis=1)
        dR = dR / np.sqrt(np.sum(dR ** 2, axis=0))
        dR = dR * self.res

        numConsidered = 0
        for _ in range(2000):
            mindisttogoal = 1000000
            node = None
            numConsidered += 1
            for k in range(numofdirs):

                next = path[-1] + dR[:, k]
                # Check if this direction is valid
                if (next[0] < self.boundary[0, 0] or next[0] > self.boundary[0, 3] or \
                        next[1] < self.boundary[0, 1] or next[1] > self.boundary[0, 4] or \
                        next[2] < self.boundary[0, 2] or next[2] > self.boundary[0, 5]):
                    continue

                valid = True
                for k in range(self.blocks.shape[0]):
                    if (next[0] > self.blocks[k, 0] and next[0] < self.blocks[k, 3] and \
                            next[1] > self.blocks[k, 1] and next[1] < self.blocks[k, 4] and \
                            next[2] > self.blocks[k, 2] and next[2] < self.blocks[k, 5]):
                        valid = False
                        break
                if not valid:
                    continue

                # Update next node
                disttogoal = sum((next - goal) ** 2)
                if (disttogoal < mindisttogoal):
                    mindisttogoal = disttogoal
                    node = next

            if node is None:
                break

            path.append(node)

            # Check if done
            if sum((path[-1] - goal) ** 2) <= 0.1:
                break

        return np.array(path), numConsidered

class CollisionPlanner:
    __slots__ = ['boundary', 'blocks', 'res', 'blk_list']

    def __init__(self, boundary, blocks, res):
        self.boundary = boundary
        self.blocks = blocks
        self.res = res
        # create a list of AABBs for pyrr
        self.blk_list = list()
        for k in range(blocks.shape[0]):
            mi = blocks[k, :3].reshape(1, 3)
            ma = blocks[k, 3:6].reshape(1, 3)
            ab = np.vstack((mi, ma))
            abblk = aabb.create_from_points(ab)
            # print(mi.shape, ma.shape, ab.shape, abblk.shape)
            self.blk_list.append(abblk)

    def intersect(self, start, end, k):
        fst = 0
        fet = 1
        box = self.blk_list[k]
        # print(box.shape)
        bomin = aabb.minimum(box)
        bomax = aabb.maximum(box)
        # print(bomin.shape, bomax.shape)
        for i in range(3):
            bmin = bomin[i]
            bmax = bomax[i]
            si = start[i]
            ei = end[i]
            if si < ei:
                if si > bmax or ei < bmin:
                    return False
                di = ei - si
                st = (bmin - si) / di if si < bmin else 0
                et = (bmax - si) / di if ei > bmax else 1
            else:
                if ei > bmax or si < bmin:
                    return False
                di = ei - si
                st = (bmax - si) / di if si > bmax else 0
                et = (bmin - si) / di if ei < bmin else 1
            if st > fst:
                fst = st
            if et < fet:
                fet = et
            if fet < fst:
                return False
        return True

    def ray_intersect_aabb(self, ray, k):
        """Calculates the intersection point of a ray and an AABB
        :param numpy.array ray1: The ray to check.
        :param numpy.array aabb: The Axis-Aligned Bounding Box to check against.
        :rtype: numpy.array
        :return: Returns a vector if an intersection occurs.
            Returns None if no intersection occurs.
        """
        """
        http://gamedev.stackexchange.com/questions/18436/most-efficient-aabb-vs-ray-collision-algorithms
        """
        # this is basically "numpy.divide( 1.0, ray[ 1 ] )"
        # except we're trying to avoid a divide by zero warning
        # so where the ray direction value is 0.0, just use infinity
        # which is what we want anyway
        direction = ray[1]
        dir_fraction = np.empty(3, dtype=ray.dtype)
        dir_fraction[direction == 0.0] = 1e8
        dir_fraction[direction != 0.0] = np.divide(1.0, direction[direction != 0.0])

        aabb = self.blk_list[k]
        t1 = (aabb[0, 0] - ray[0, 0]) * dir_fraction[0]
        t2 = (aabb[1, 0] - ray[0, 0]) * dir_fraction[0]
        t3 = (aabb[0, 1] - ray[0, 1]) * dir_fraction[1]
        t4 = (aabb[1, 1] - ray[0, 1]) * dir_fraction[1]
        t5 = (aabb[0, 2] - ray[0, 2]) * dir_fraction[2]
        t6 = (aabb[1, 2] - ray[0, 2]) * dir_fraction[2]

        tmin = max(min(t1, t2), min(t3, t4), min(t5, t6))
        tmax = min(max(t1, t2), max(t3, t4), max(t5, t6))

        # if tmax < 0, ray (line) is intersecting AABB
        # but the whole AABB is behind the ray start
        if tmax < 0:
            return None

        # if tmin > tmax, ray doesn't intersect AABB
        if tmin > tmax:
            return None

        # t is the distance from the ray point
        # to intersection

        t = min(x for x in [tmin, tmax] if x >= 0)
        point = ray[0] + (ray[1] * t)
        return point

    def plan(self, start, goal):
        path = [start]
        numofdirs = 26
        [dX, dY, dZ] = np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])
        dR = np.vstack((dX.flatten(), dY.flatten(), dZ.flatten()))
        dR = np.delete(dR, 13, axis=1)
        dR = dR / np.sqrt(np.sum(dR ** 2, axis=0))
        dR = dR * self.res

        numConsidered = 0
        for _ in range(2000):
            mindisttogoal = 1000000
            node = None
            numConsidered += 1
            prev = path[-1]
            for j in range(numofdirs):
                next = path[-1] + dR[:, j]
                # Check if this direction is within boundaries
                if (next[0] < self.boundary[0, 0] or next[0] > self.boundary[0, 3] or \
                        next[1] < self.boundary[0, 1] or next[1] > self.boundary[0, 4] or \
                        next[2] < self.boundary[0, 2] or next[2] > self.boundary[0, 5]):
                    continue

                valid = True
                # loop through all obstacles
                for k in range(self.blocks.shape[0]):
                    # check if next node is within any one of these blocks
                    if (next[0] > self.blocks[k, 0] and next[0] < self.blocks[k, 3] and \
                            next[1] > self.blocks[k, 1] and next[1] < self.blocks[k, 4] and \
                            next[2] > self.blocks[k, 2] and next[2] < self.blocks[k, 5]):
                        valid = False
                        break
                    # check for intersection
                    # check if ray from node to next intersects AABB
                    # se = line.create_from_points(prev, next)
                    # es = line.create_from_points(next, prev)
                    # rf = ray.create_from_line(se)
                    # rb = ray.create_from_line(es)
                    # rfi = self.ray_intersect_aabb(rf, k)
                    # rbi = self.ray_intersect_aabb(rb, k)
                    rfi = self.intersect(prev, next, k)
                    rbi = self.intersect(next, prev, k)
                    # check if collision occurred in between these 2 points
                    if rfi and rbi:
                        # if rfi is not None and rbi is not None:
                        valid = False
                        break
                if not valid:
                    continue

                # Update next node
                disttogoal = sum((next - goal) ** 2)
                if (disttogoal < mindisttogoal):
                    mindisttogoal = disttogoal
                    node = next

            if node is None:
                break

            path.append(node)

            # Check if done
            if sum((path[-1] - goal) ** 2) <= 0.1:
                break

        return np.array(path), numConsidered

def thas(t):
    return t

class AStarPlanner:
    __slots__ = ['boundary', 'blocks', 'res', 'blk_list']

    def __init__(self, boundary, blocks, res):
        self.boundary = boundary
        self.blocks = blocks
        self.res = res
        # create a list of AABBs for pyrr
        self.blk_list = list()
        for k in range(blocks.shape[0]):
            mi = blocks[k, :3].reshape(1, 3)
            ma = blocks[k, 3:6].reshape(1, 3)
            ab = np.vstack((mi, ma))
            abblk = aabb.create_from_points(ab)
            # print(mi.shape, ma.shape, ab.shape, abblk.shape)
            self.blk_list.append(abblk)

    def l2_dist(self, s, e):
        """
        return the L2 distance between 2 points
        :param s: node 1
        :param e: node 2
        :return: L2 distance between s and e
        """
        return np.sqrt(sum((s - e) ** 2))

    def l_inf_dist(self, s, e):
        """
        return the L infinity distance between 2 points
        :param s: node 1
        :param e: node 2
        :return: l_inf distance between s and e
        """
        return np.amax(s - e)

    def heuristic(self, n, goal):
        """
        return the heuristic between current node and goal
        :param n: current node
        :param goal: goal node
        :return: heuristic function
        """
        # function can either be L2 or L_inf
        return self.l2_dist(n, goal)

    def intersect(self, start, end, k):
        fst = 0
        fet = 1
        box = self.blk_list[k]
        # print(box.shape)
        bomin = aabb.minimum(box)
        bomax = aabb.maximum(box)
        # print(bomin.shape, bomax.shape)
        for i in range(3):
            bmin = bomin[i]
            bmax = bomax[i]
            si = start[i]
            ei = end[i]
            if si < ei:
                if si > bmax or ei < bmin:
                    return False
                di = ei - si
                st = (bmin - si) / di if si < bmin else 0
                et = (bmax - si) / di if ei > bmax else 1
            else:
                if ei > bmax or si < bmin:
                    return False
                di = ei - si
                st = (bmax - si) / di if si > bmax else 0
                et = (bmin - si) / di if ei < bmin else 1
            if st > fst:
                fst = st
            if et < fet:
                fet = et
            if fet < fst:
                return False
        return True

    def reachable_goal(self, n, g):
        # check if goal is validly reachable
        for k in range(self.blocks.shape[0]):
            rfi = self.intersect(n, g, k)
            rbi = self.intersect(g, n, k)
            # check if collision occurred in between these 2 points
            if rfi and rbi:
                return False
        return True

    def plan(self, start, goal):
        numofdirs = 26
        [dX, dY, dZ] = np.meshgrid([-1, 0, 1], [-1, 0, 1], [-1, 0, 1])
        dR = np.vstack((dX.flatten(), dY.flatten(), dZ.flatten()))
        dC = np.delete(dR, 13, axis=1)
        dR = dC / np.sqrt(np.sum(dC ** 2, axis=0))
        dR = dR * self.res
        thresh = (self.res / 2.0) * np.sqrt(3)
        INF = 9999999

        numConsidered = 0
        # OPEN is a priority queue
        pq = pqdict(key=lambda x: x['f'])
        g_tau = INF
        # g value dict
        gs = dict()
        # parent dict
        parents = dict()

        # insert start node
        # s = str((0, 0, 0))
        s = (0, 0, 0)
        hs = thas(s)
        dat = dict()
        dat['ind'] = s
        dat['coord'] = start
        dat['g'] = 0.0
        gs[hs] = 0.0
        dat['h'] = self.heuristic(start, goal)
        dat['f'] = dat['g'] + dat['h']
        pd = (None, None)
        parents[hs] = pd
        pq[hs] = dat
        print(pq, len(pq))
        fin_i = None
        fin_g = None

        # while OPEN is not empty
        while len(pq) is not 0:
            i = pq.popitem()
            # print(i)
            numConsidered += 1
            i_dat = i[1]
            i_ind = np.array(i_dat['ind'])
            # i_ind = np.array(i[0])
            i_coord = np.array(i_dat['coord'])
            # i_to_g = np.sqrt(sum((i_coord - goal) ** 2))
            i_to_g = self.l2_dist(i_coord, goal)
            if i_to_g <= thresh:
                if self.reachable_goal(i_coord, goal):
                    print("Goal reached!!")
                    fin_i = tuple(i_ind)
                    fin_g = i_coord
                    assert ((g_tau - i_dat['g'] - i_to_g) <= 1e-4)
                    break
            for j in range(numofdirs):
                # next = i_coord + dR[:, j]
                next = i_coord + dC[:, j] * self.res
                next_ind = i_ind + dC[:, j]
                # Check if this direction is within boundaries
                if (next[0] < self.boundary[0, 0] or next[0] > self.boundary[0, 3] or \
                        next[1] < self.boundary[0, 1] or next[1] > self.boundary[0, 4] or \
                        next[2] < self.boundary[0, 2] or next[2] > self.boundary[0, 5]):
                    continue

                valid = True
                # loop through all obstacles
                for k in range(self.blocks.shape[0]):
                    # check if next node is within any one of these blocks
                    if (next[0] > self.blocks[k, 0] and next[0] < self.blocks[k, 3] and \
                            next[1] > self.blocks[k, 1] and next[1] < self.blocks[k, 4] and \
                            next[2] > self.blocks[k, 2] and next[2] < self.blocks[k, 5]):
                        valid = False
                        break
                    # check for intersection
                    rfi = self.intersect(i_coord, next, k)
                    rbi = self.intersect(next, i_coord, k)
                    # check if collision occurred in between these 2 points
                    if rfi and rbi:
                        valid = False
                        break
                if not valid:
                    continue

                # cs = np.array([2.07313218, 3.56684973, 0.75795371])
                # ce = np.array([2.2028887, 0.95360107, 0.75795371])
                # the potential g value
                upg = i_dat['g'] + self.l2_dist(i_coord, next)
                # heuristic for next node
                hj = self.heuristic(next, goal)
                ni = thas(tuple(next_ind))
                # ni = tuple(next_ind)
                if ni not in gs.keys():
                    gj = INF
                else:
                    gj = gs[ni]
                # checking the OPEN condition
                if upg < gj and upg + hj < g_tau:
                    # if abs(sum(i_coord - cs)) <= 1e-4 \
                    #         and abs(sum(next - ce)) <= 1e-4:
                    #     print("MONZA - Now it does NOT detect a collision!!")
                    gs[ni] = upg
                    assert (self.l2_dist(i_coord, next) <= 1e-4 + 2.0 * thresh)
                    parents[ni] = (tuple(i_ind), tuple(i_coord))
                    next_dat = dict()
                    next_dat['coord'] = next
                    next_dat['ind'] = tuple(next_ind)
                    next_dat['g'] = upg
                    next_dat['h'] = hj
                    next_dat['f'] = upg + hj
                    pq[ni] = next_dat
                    # check if newly inserted is goal
                    # if yes, update g_tau
                    if self.l2_dist(next, goal) <= 1e-4 + thresh:
                        if self.reachable_goal(next, goal):
                            g_tau = min(g_tau, upg + hj)

        # if no path exists, i.e., goal was not reached
        if fin_i is None:
            print("Goal was NOT reachable!")
            path = [start]
            print(type(path))
            return np.array(path), numConsidered

        # if path exists recover it
        path = [goal]
        print(type(path))
        path.append(fin_g)
        cur = thas(fin_i)
        # cur = fin_i
        prev = parents[cur][1]
        while prev is not None:
            print(len(path))
            print(type(prev), prev)
            if self.l2_dist(np.array(prev), path[-1]) > 2.0*thresh + 1e-4:
                print("NOOOO!")
                print("prev: {}; cur: {}\nparents[cur]: {}; hash(par): {}".format(prev, cur, parents[cur],
                                                                                thas(parents[cur][0])))
            assert (self.l2_dist(np.array(prev), path[-1]) <= 1e-4 + 2.0 * thresh)
            path.append(np.array(prev))
            cur = thas(parents[cur][0])
            # cur = parents[cur][0]
            prev = parents[cur][1]
        path.reverse()
        path = np.array(path)
        print(type(path[0]), type(start))
        print(path[0].shape, start.shape)
        if id(cur) is not id(thas(s)):
            print("cur: {}; thas(s): {}".format(cur, thas(s)))
            print("id(cur): {}; id(s): {}".format(id(cur), id(thas(s))))
        assert (hash(cur) == hash(thas(s)))
        assert (abs(sum(path[0] - start)) < 1e-4)
        return np.array(path), numConsidered


text:
import pandas as pd
import numpy as np

def get_tensile_sensor_names():

    tensile_sensor_names = [
        "Bolt_1_Tensile",
        "Bolt_2_Tensile",
        "Bolt_3_Tensile",
        "Bolt_4_Tensile",
        "Bolt_5_Tensile",
        "Bolt_6_Tensile",
    ]

    return tensile_sensor_names

def get_torsion_sensor_names():

    torsion_sensor_names = [
        "Bolt_1_Torsion",
        "Bolt_2_Torsion",
        "Bolt_3_Torsion",
        "Bolt_4_Torsion",
        "Bolt_5_Torsion",
        "Bolt_6_Torsion",
    ]

    return torsion_sensor_names

def get_operating_signal_names():

    operating_signal_names = [
        "Unit_4_Power",
        "Unit_4_Reactive Power",
        "Turbine_Guide Vane Opening",
        "Turbine_Pressure Drafttube",
        "Turbine_Pressure Spiral Casing",
        "Turbine_Rotational Speed",
    ]

    return operating_signal_names

def get_running_state(df, minutes_of_masking: int):

    steady_df = df.copy()

    start_idx = steady_df.operation_mode == "start"
    steady_df = steady_df.drop(columns=["operation_mode"])

    start_times = steady_df.index[start_idx]
    minute_mask = pd.to_timedelta(minutes_of_masking, 'm')

    for start_time in start_times:
        end_time = start_time + minute_mask
        steady_df.loc[start_time:end_time, :] = np.nan

    return steady_df

def get_transient_state(df, minutes_of_masking: int):

    transient_df = df.copy()
    start_idx = transient_df.operation_mode == "start"

    start_times = transient_df.index[start_idx]
    minute_mask = pd.to_timedelta(minutes_of_masking, "m")

    for start_time in start_times:
        end_time = start_time + minute_mask
        transient_df.loc[start_time:end_time, "operation_mode"] = "start"

    transient_df.loc[transient_df.operation_mode!="start",:] = np.nan
    transient_df = transient_df.drop(columns=['operation_mode'])

    return transient_df

def load_training_data(
    full: bool, steady: bool, minutes_of_masking: int
) -> pd.DataFrame:
    """Get training dataset

    Args:
        full (bool): include all data avialable
        steady (bool): remove every observation after a start operation_mode
        minutes_of_masking (int): number of minutes to mask
    Returns:
        pd.DataFrame: dataframe with training data, either steady or transient state
    """

    df = pd.read_parquet("../data/input_dataset-2.parquet")

    if full:
        df_1 = pd.read_parquet("../data/input_dataset-1.parquet")
        df = pd.concat([df_1, df])

    df = df.rename(columns={"mode": "operation_mode"})

    if steady:
        df = get_running_state(df, minutes_of_masking=minutes_of_masking)
    if not steady:
        df = get_transient_state(df, minutes_of_masking=minutes_of_masking)

    df = df.reindex()

    return df

def load_test_data():

    a = 1


text:
import logging

import numpy as np
from sklearn.cluster import MiniBatchKMeans
from sklearn.utils import check_random_state

from csrank.discretechoice.likelihoods import create_weight_dictionary
from csrank.discretechoice.likelihoods import fit_pymc3_model
from csrank.learner import Learner
import csrank.numpy_util as npu
import csrank.theano_util as ttu
from csrank.util import print_dictionary
from .discrete_choice import DiscreteObjectChooser
from .likelihoods import likelihood_dict
from .likelihoods import LogLikelihood

try:
    import pymc3 as pm
    from pymc3.variational.callbacks import CheckParametersConvergence
except ImportError:
    from csrank.util import MissingExtraError

    raise MissingExtraError("pymc3", "probabilistic")

try:
    import theano
    from theano import tensor as tt
except ImportError:
    from csrank.util import MissingExtraError

    raise MissingExtraError("theano", "probabilistic")

logger = logging.getLogger(__name__)

class NestedLogitModel(DiscreteObjectChooser, Learner):
    def __init__(
        self,
        n_nests=None,
        loss_function="",
        regularization="l1",
        alpha=1e-2,
        random_state=None,
        **kwd,
    ):
        """
        Create an instance of the Nested Logit model for learning the discrete choice function. This model divides
        objects into disjoint subsets called nests,such that the objects which are similar to each other are in same
        nest. This model structure is 1-layer of hierarchy and the :math:`\\lambda` for each nest :math:`B_k` signifies
        the degree of independence and  :math:`1-\\lambda` signifies the correlations between the object in it. We
        learn two weight vectors and the  :math:`\\lambda s`.

        The probability of choosing an object :math:`x_i` from the given query set :math:`Q` is defined by product
        of choosing the nest in which :math:`x_i` exists and then choosing the the object from the nest.

        .. math::

            P(x_i \\lvert Q) = P_i = P_{i \\lvert B_k} P_{B_k} \\enspace ,

        The discrete choice for the given query set :math:`Q` is defined as:

        .. math::

            dc(Q) := \\operatorname{argmax}_{x_i \\in Q }  \\; P(x_i \\lvert Q)

        Parameters
        ----------
        n_nests : int range : [2,n_objects/2]
            The number of nests/subsets in which the objects are divided.
            This may not surpass half the amount of objects this model will
            be trained on.
        loss_function : string , {‘categorical_crossentropy’, ‘binary_crossentropy’, ’categorical_hinge’}
            Loss function to be used for the discrete choice decision from the query set
        regularization : string, {‘l1’, ‘l2’}, string
           Regularizer function (L1 or L2) applied to the `kernel` weights matrix
        alpha: float (range : [0,1])
            The lower bound of the correlations between the objects in a nest
        random_state : int or object
            Numpy random state
        **kwargs
            Keyword arguments for the algorithms

        References
        ----------
            [1] Kenneth E Train. „Discrete choice methods with simulation“. In: Cambridge university press, 2009. Chap GEV, pp. 87–111.

            [2] Kenneth Train. Qualitative choice analysis. Cambridge, MA: MIT Press, 1986

            [3] Kenneth Train and Daniel McFadden. „The goods/leisure tradeoff and disaggregate work trip mode choice models“. In: Transportation research 12.5 (1978), pp. 349–353
        """
        self.n_nests = n_nests
        self.alpha = alpha
        self.random_state = random_state
        self.loss_function = loss_function
        known_regularization_functions = {"l1", "l2"}
        if regularization not in known_regularization_functions:
            raise ValueError(
                f"Regularization function {regularization} is unknown. Must be one of {known_regularization_functions}"
            )
        self.regularization = regularization

    @property
    def model_configuration(self):
        """
            Constructs the dictionary containing the priors for the weight vectors for the model according to the
            regularization function. The parameters are:
                * **weights** : Weights to evaluates the utility of the objects
                * **weights_k** : Weights to evaluates the utility of the nests

            For ``l1`` regularization the priors are:

            .. math::

                \\text{mu}_w \\sim \\text{Normal}(\\text{mu}=0, \\text{sd}=5.0) \\\\
                \\text{b}_w \\sim \\text{HalfCauchy}(\\beta=1.0) \\\\
                \\text{weights} \\sim \\text{Laplace}(\\text{mu}=\\text{mu}_w, \\text{b}=\\text{b}_w)

            For ``l2`` regularization the priors are:

            .. math::

                \\text{mu}_w \\sim \\text{Normal}(\\text{mu}=0, \\text{sd}=5.0) \\\\
                \\text{sd}_w \\sim \\text{HalfCauchy}(\\beta=1.0) \\\\
                \\text{weights} \\sim \\text{Normal}(\\text{mu}=\\text{mu}_w, \\text{sd}=\\text{sd}_w)

            Returns
            -------
                configuration : dict
                    Dictionary containing the priors applies on the weights
        """
        if not hasattr(self, "config_"):
            if self.regularization == "l2":
                weight = pm.Normal
                prior = "sd"
            elif self.regularization == "l1":
                weight = pm.Laplace
                prior = "b"
            self.config_ = {
                "weights": [
                    weight,
                    {
                        "mu": (pm.Normal, {"mu": 0, "sd": 5}),
                        prior: (pm.HalfCauchy, {"beta": 1}),
                    },
                ],
                "weights_k": [
                    weight,
                    {
                        "mu": (pm.Normal, {"mu": 0, "sd": 5}),
                        prior: (pm.HalfCauchy, {"beta": 1}),
                    },
                ],
            }
            logger.info(
                "Creating model with config {}".format(print_dictionary(self.config_))
            )
        return self.config_

    def create_nests(self, X):
        """
        For allocating the objects to different nests we use the clustering algorithm with number of clusters
        :math:`k` and allocate the similar objects in query set :math:`Q`.

        Parameters
        ----------
        X : numpy array
            (n_instances, n_objects, n_features)
            Feature vectors of the objects in the query sets

        Returns
        -------
        Yn : numpy array
            (n_instances, n_objects) Values for each object implying the nest it belongs to. For example for :math:`2` nests the value 0 implies that object is allocated to nest 1 and value 1 implies it is allocated to nest 2.

        """
        self.random_state_ = self.random_state_
        n, n_obj, n_dim = X.shape
        objects = X.reshape(n * n_obj, n_dim)
        if not hasattr(self, "cluster_model_"):
            self.cluster_model_ = MiniBatchKMeans(
                n_clusters=self.n_nests, random_state=self.random_state_
            ).fit(objects)
            self.features_nests_ = self.cluster_model_.cluster_centers_
            prediction = self.cluster_model_.labels_
        else:
            prediction = self.cluster_model_.predict(objects)
        Yn = []
        for i in np.arange(0, n * n_obj, step=n_obj):
            nest_ids = prediction[i : i + n_obj]
            Yn.append(nest_ids)
        Yn = np.array(Yn)
        return Yn

    def _eval_utility(self, weights):
        utility = tt.zeros(tuple(self.y_nests_.shape))
        for i in range(self.n_nests):
            rows, cols = tt.eq(self.y_nests_, i).nonzero()
            utility = tt.set_subtensor(
                utility[rows, cols], tt.dot(self.Xt_[rows, cols], weights[i])
            )
        return utility

    def get_probabilities(self, utility, lambda_k, utility_k):
        """
            This method calculates the probability of choosing an object from the query set using the following parameters of the model which are used:

                * **weights** (:math:`w`): Weights to get the utility of the object :math:`Y_i = U(x_i) = w \\cdot x_i`
                * **weights_k** (:math:`w_k`): Weights to get the utility of the next  :math:`W_k = U_k(x) = w_k \\cdot c_k`, where :math:`c_k` is the center of the object space of nest :math:`B_k`
                * **lambda_k** (:math:`\\lambda_k`): Lambda is the measure of independence amongst the obejcts in the nest :math:`B_k`

            The probability of choosing the object  :math:`x_i` from the query set :math:`Q`:

            .. math::
                P_i = \\frac{\\boldsymbol{e}^{ ^{Y_i} /_{\\lambda_k}}}{\\sum_{j \\in B_k} \\boldsymbol{e}^{^{Y_j} /_{\\lambda_k}}} \\frac {\\boldsymbol{e}^{W_k + \\lambda_k I_k}} {\\sum_{\\ell = 1}^{K} \\boldsymbol{e}^{ W_{\\ell } + \\lambda_{\\ell} I_{\\ell}}} \\quad i \\in B_k  \\enspace , \\\\
                where,\\enspace I_k = \\ln \\sum_{ j \\in B_k} \\boldsymbol{e}^{^{Y_j} /_{\\lambda_k}}

            Parameters
            ----------
            utility : theano tensor
                (n_instances, n_objects)
                Utility :math:`Y_i` of the objects :math:`x_i \\in Q` in the query sets
            lambda_k : theano tensor (range : [alpha, 1.0])
                (n_nests)
                Measure of independence amongst the obejcts in each nests
            utility_k : theano tensor
                (n_instances, n_nests)
                Utilities of the nests :math:`B_k \\in \\mathcal{B}`

            Returns
            -------
            p : theano tensor
                (n_instances, n_objects)
                Choice probabilities :math:`P_i` of the objects :math:`x_i \\in Q` in the query sets

        """
        n_instances, n_objects = self.y_nests_.shape
        pni_k = tt.zeros((n_instances, n_objects))
        ivm = tt.zeros((n_instances, self.n_nests))
        for i in range(self.n_nests):
            rows, cols = tt.neq(self.y_nests_, i).nonzero()
            sub_tensor = tt.set_subtensor(utility[rows, cols], -1e50)
            ink = ttu.logsumexp(sub_tensor)
            rows, cols = tt.eq(self.y_nests_, i).nonzero()
            pni_k = tt.set_subtensor(
                pni_k[rows, cols], tt.exp(sub_tensor - ink)[rows, cols]
            )
            ivm = tt.set_subtensor(ivm[:, i], lambda_k[i] * ink[:, 0] + utility_k[i])
        pk = tt.exp(ivm - ttu.logsumexp(ivm))
        pn_k = tt.zeros((n_instances, n_objects))
        for i in range(self.n_nests):
            rows, cols = tt.eq(self.y_nests_, i).nonzero()
            p = tt.ones((n_instances, n_objects)) * pk[:, i][:, None]
            pn_k = tt.set_subtensor(pn_k[rows, cols], p[rows, cols])
        p = pni_k * pn_k
        return p

    def _eval_utility_np(self, x_t, y_nests, weights):
        utility = np.zeros(tuple(y_nests.shape))
        for i in range(self.n_nests):
            rows, cols = np.where(y_nests == i)
            utility[rows, cols] = np.dot(x_t[rows, cols], weights[i])
        return utility

    def _get_probabilities_np(self, Y_n, utility, lambda_k, utility_k):
        n_instances, n_objects = Y_n.shape
        pni_k = np.zeros((n_instances, n_objects))
        ivm = np.zeros((n_instances, self.n_nests))
        for i in range(self.n_nests):
            sub_tensor = np.copy(utility)
            sub_tensor[np.where(Y_n != i)] = -1e50
            ink = npu.logsumexp(sub_tensor)
            pni_k[np.where(Y_n == i)] = np.exp(sub_tensor - ink)[np.where(Y_n == i)]
            ivm[:, i] = lambda_k[i] * ink[:, 0] + utility_k[i]
        pk = np.exp(ivm - npu.logsumexp(ivm))
        pn_k = np.zeros((n_instances, n_objects))
        for i in range(self.n_nests):
            rows, cols = np.where(Y_n == i)
            p = np.ones((n_instances, n_objects)) * pk[:, i][:, None]
            pn_k[rows, cols] = p[rows, cols]
        p = pni_k * pn_k
        return p

    def construct_model(self, X, Y):
        """
        Constructs the nested logit model by applying priors on weight vectors **weights** and **weights_k** as per
        :meth:`model_configuration`. Then we apply a uniform prior to the :math:`\\lambda s`, i.e.
        :math:`\\lambda s \\sim Uniform(\\text{alpha}, 1.0)`.The probability of choosing the object :math:`x_i` from
        the query set :math:`Q = \\{x_1, \\ldots ,x_n\\}` is evaluated in :meth:`get_probabilities`.

        Parameters
        ----------
        X : numpy array
            (n_instances, n_objects, n_features)
            Feature vectors of the objects
        Y : numpy array
            (n_instances, n_objects)
            Preferences in the form of discrete choices for given objects

        Returns
        -------
         model : pymc3 Model :class:`pm.Model`
        """
        self.loss_function_ = likelihood_dict.get(self.loss_function, None)
        self.threshold_ = 5e6
        self.trace_ = None
        self.trace_vi_ = None
        if np.prod(X.shape) > self.threshold_:
            upper_bound = int(self.threshold_ / np.prod(X.shape[1:]))
            indices = self.random_state_.choice(X.shape[0], upper_bound, replace=False)
            X = X[indices, :, :]
            Y = Y[indices, :]
        logger.info("Train Set instances {} objects {} features {}".format(*X.shape))
        y_nests = self.create_nests(X)
        with pm.Model() as self.model:
            self.Xt_ = theano.shared(X)
            self.Yt_ = theano.shared(Y)
            self.y_nests_ = theano.shared(y_nests)
            shapes = {
                "weights": self.n_object_features_fit_,
                "weights_k": self.n_object_features_fit_,
            }

            weights_dict = create_weight_dictionary(self.model_configuration, shapes)
            lambda_k = pm.Uniform("lambda_k", self.alpha, 1.0, shape=self.n_nests)
            weights = weights_dict["weights"] / lambda_k[:, None]
            utility = self._eval_utility(weights)
            utility_k = tt.dot(self.features_nests_, weights_dict["weights_k"])
            self.p_ = self.get_probabilities(utility, lambda_k, utility_k)

            LogLikelihood(
                "yl", loss_func=self.loss_function_, p=self.p_, observed=self.Yt_
            )
        logger.info("Model construction completed")

    def fit(
        self,
        X,
        Y,
        sampler="variational",
        tune=500,
        draws=500,
        vi_params={
            "n": 20000,
            "method": "advi",
            "callbacks": [CheckParametersConvergence()],
        },
        **kwargs,
    ):
        """
        Fit a nested logit model on the provided set of queries X and choices Y of those objects. The provided
        queries and corresponding preferences are of a fixed size (numpy arrays). For learning this network the
        categorical cross entropy loss function for each object :math:`x_i \\in Q` is defined as:

        .. math::

            C_{i} =  -y(i)\\log(P_i) \\enspace,

        where :math:`y` is ground-truth discrete choice vector of the objects in the given query set :math:`Q`.
        The value :math:`y(i) = 1` if object :math:`x_i` is chosen else :math:`y(i) = 0`.

        Parameters
        ----------
        X : numpy array (n_instances, n_objects, n_features)
            Feature vectors of the objects
        Y : numpy array (n_instances, n_objects)
            Choices for given objects in the query
        sampler : {‘variational’, ‘metropolis’, ‘nuts’}, string
            The sampler used to estimate the posterior mean and mass matrix from the trace

                * **variational** : Run inference methods to estimate posterior mean and diagonal mass matrix
                * **metropolis** : Use the MAP as starting point and Metropolis-Hastings sampler
                * **nuts** : Use the No-U-Turn sampler
        vi_params : dict
            The parameters for the **variational** inference method
        draws : int
            The number of samples to draw. Defaults to 500. The number of tuned samples are discarded by default
        tune : int
            Number of iterations to tune, defaults to 500. Ignored when using 'SMC'. Samplers adjust
            the step sizes, scalings or similar during tuning. Tuning samples will be drawn in addition
            to the number specified in the `draws` argument, and will be discarded unless
            `discard_tuned_samples` is set to False.
        **kwargs :
            Keyword arguments for the fit function of :meth:`pymc3.fit`or :meth:`pymc3.sample`
        """
        self._pre_fit()
        _n_instances, self.n_objects_fit_, self.n_object_features_fit_ = X.shape
        if self.n_nests is None:
            self.n_nests = int(self.n_objects_fit_ / 2)
        self.random_state_ = check_random_state(self.random_state)
        self.construct_model(X, Y)
        fit_pymc3_model(self, sampler, draws, tune, vi_params, **kwargs)
        return self

    def _predict_scores_fixed(self, X, **kwargs):
        y_nests = self.create_nests(X)
        mean_trace = dict(pm.summary(self.trace_)["mean"])
        weights = np.array(
            [
                mean_trace["weights[{}]".format(i)]
                for i in range(self.n_object_features_fit_)
            ]
        )
        weights_k = np.array(
            [
                mean_trace["weights_k[{}]".format(i)]
                for i in range(self.n_object_features_fit_)
            ]
        )
        lambda_k = np.array(
            [mean_trace["lambda_k[{}]".format(i)] for i in range(self.n_nests)]
        )
        weights = weights / lambda_k[:, None]
        utility_k = np.dot(self.features_nests_, weights_k)
        utility = self._eval_utility_np(X, y_nests, weights)
        scores = self._get_probabilities_np(y_nests, utility, lambda_k, utility_k)
        return scores

max_stars_repo_name: kiudee/cs-ranking

max_issues_repo_name: kiudee/cs-ranking

max_forks_repo_name: kiudee/cs-ranking

reason: import numpy,import theano,from theano,import pymc3,from pymc3


text:
import os
import sys
from sklearn.datasets import fetch_mldata
import numpy as np

def digit_to_vec(x):
    res = np.zeros(10, dtype=np.float32)
    res[int(x)] = 1
    return res

def get_dataset_ori():
    mnist = fetch_mldata('MNIST original')
    return mnist.data, mnist.target

def get_dataset():
    X, y = get_dataset_ori()
    X = (X / 255.0).astype(np.float32)
    y = np.array([digit_to_vec(n) for n in y])
    return X, y

def save_bin_dataset(path):
    if os.path.isfile(path):
        return

    print('Generating MNIST binary file')
    X, y = get_dataset_ori()
    with open(path, 'wb') as f:
        for i in range(len(X)):
            f.write(bytes(X[i].tolist()))
            f.write(bytes([int(y[i])]))
    print('Saved to ' + path)

if __name__ == '__main__':
    save_bin_dataset(sys.argv[1])


text:
import numpy as np

from snake import SnakeEnv, ModelFreeAgent, TableAgent, eval_game
import gym
from policy_iter import PolicyIteration
from monte_carlo import MonteCarlo,timer

class SARSA(object):
    def __init__(self, epsilon=0.0):
        self.epsilon = epsilon

    def sarsa_eval(self, agent, env):
        #sarsa
        state = env.reset()
        prev_state = -1
        prev_act = -1
        while True:
            act = agent.play(state, self.epsilon)
            next_state, reward, terminate, _ = env.step(act)
            if prev_act != -1:
                return_val = reward + agent.gamma * (0 if terminate else agent.value_q[state][act])
                agent.value_n[prev_state][prev_act] += 1
                agent.value_q[prev_state][prev_act] += (return_val - \
                    agent.value_q[prev_state][prev_act]) / \
                    agent.value_n[prev_state][prev_act]

            prev_act = act
            prev_state = state
            state = next_state

            if terminate:
                break

    def policy_improve(self, agent):
        new_policy = np.zeros_like(agent.pi)
        for i in range(1, agent.s_len):
            new_policy[i] = np.argmax(agent.value_q[i,:])
        if np.all(np.equal(new_policy, agent.pi)):
            return False
        else:
            agent.pi = new_policy
            return True

    # monte carlo
    def sarsa(self, agent, env):
        for i in range(10):
            for j in range(2000):
                self.sarsa_eval(agent, env)
            self.policy_improve(agent)

def monte_carlo_demo():
    np.random.seed(101)
    env = SnakeEnv(10, [3,6])
    agent = ModelFreeAgent(env)
    mc = MonteCarlo(0.5)
    with timer('Timer Monte Carlo Iter'):
        mc.monte_carlo_opt(agent, env)
    print 'return_pi={}'.format(eval_game(env,agent))
    print agent.pi

    np.random.seed(101)
    agent2 = TableAgent(env)
    pi_algo = PolicyIteration()
    with timer('Timer PolicyIter'):
        pi_algo.policy_iteration(agent2)
    print 'return_pi={}'.format(eval_game(env,agent2))
    print agent2.pi

    np.random.seed(101)
    agent3 = ModelFreeAgent(env)
    mc = SARSA(0.5)
    with timer('Timer Monte Carlo Iter'):
        mc.sarsa(agent3, env)
    print 'return_pi={}'.format(eval_game(env,agent3))
    print agent3.pi

if __name__ == '__main__':
    monte_carlo_demo()

max_stars_repo_name: jayhenry/rlpyt

max_issues_repo_name: jayhenry/rlpyt

max_forks_repo_name: jayhenry/rlpyt