text:
#!/usr/bin/env python
"""
sen1_ard_gamma - tools for Sentinel-1 GRD processing using Gamma

"""
# This file is part of 'sen1_ard_gamma'
# A set of tools to produce Sentinel-1 ARD using Gamma.
#
# Copyright 2019 Pete Bunting
#
# 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.
#
#
# Purpose:  Setup variables and imports across the whole module
#
# Author: Pete Bunting
# Email: pfb@aber.ac.uk
# Date: 06/06/2019
# Version: 1.0
#
# History:
# Version 1.0 - Created.

import os
import os.path
import glob
import shutil

import math

import subprocess

import numpy

import osgeo.gdal as gdal

import sen1_ard_gamma.sen1_ard_utils
import sen1_ard_gamma.gamma_dem_prep
import sen1_ard_gamma.calc_img_stats
from sen1_ard_gamma import GTIFF_CREATION_OPTS

import logging
logger = logging.getLogger(__name__)

gdal.UseExceptions()

def get_sen1_latlong_bbox(ann_sen1_xml):
    """

    :param ann_sen1_xml:
    :return:

    """
    import xml.etree.ElementTree as ET
    tree = ET.parse(ann_sen1_xml)
    root = tree.getroot()

    geo_loc_grid_tag = root.find('geolocationGrid')
    if geo_loc_grid_tag is None:
        raise Exception("Could not find the geolocationGrid tag within product tag in Annotations XML file.")

    geo_loc_grid_pt_lst_tag = geo_loc_grid_tag.find('geolocationGridPointList')
    if geo_loc_grid_pt_lst_tag is None:
        raise Exception(
            "Could not find the geolocationGridPointList tag within geolocationGrid tag in Annotations XML file.")

    min_lat = 0.0
    max_lat = 0.0
    min_lon = 0.0
    max_lon = 0.0
    first = True
    for geo_loc_grid_pt_tag in geo_loc_grid_pt_lst_tag:
        if geo_loc_grid_pt_tag.tag == 'geolocationGridPoint':
            lat_val_tag = geo_loc_grid_pt_tag.find('latitude')
            if lat_val_tag is None:
                raise Exception(
                    "Could not find the latitude tag within geolocationGridPoint tag in Annotations XML file.")
            lon_val_tag = geo_loc_grid_pt_tag.find('longitude')
            if lon_val_tag is None:
                raise Exception(
                    "Could not find the longitude tag within geolocationGridPoint tag in Annotations XML file.")
            lat_str_val = lat_val_tag.text.strip()
            lon_str_val = lon_val_tag.text.strip()

            lat_val = sen1_ard_gamma.sen1_ard_utils.str2float(lat_str_val)
            lon_val = sen1_ard_gamma.sen1_ard_utils.str2float(lon_str_val)

            if first:
                min_lat = lat_val
                max_lat = lat_val
                min_lon = lon_val
                max_lon = lon_val
                first = False
            else:
                if lat_val < min_lat:
                    min_lat = lat_val
                if lat_val > max_lat:
                    max_lat = lat_val

                if lon_val < min_lon:
                    min_lon = lon_val
                if lon_val > max_lon:
                    max_lon = lon_val

    return [min_lon, max_lon, min_lat, max_lat]

def exe_gamma_grd_ard_processing(sen1img, ann_sen1_xml, cal_sen1_xml, nos_sen1_xml, in_dem_file, demparfile,
                                 outbasename, out_dir, tmp_dir, out_res_x, out_res_y, out_proj_epsg=None,
                                 use_dem_file=True, check_in_dem_filename=False, no_dem_check=False,
                                 dem_resample_method=gdal.GRA_CubicSpline):
    """
    Function to run through a Gamma processing chain to geocode and topographically correct Sentinel-1 GRD data.

    To control the number of threads used you can use the OMP_NUM_THREADS environmental variable.

    If out_proj_epsg is None then the UTM zone for the scene centre point is identified used as the output projection.

    :param sen1img: input tiff file from SAFE file for the polarisation to be processed.
    :param ann_sen1_xml: the annotation XML file for the polarisation being processed.
    :param cal_sen1_xml: the calibration XML file for the polarisation being processed.
    :param nos_sen1_xml: the noise XML file for the polarisation being prcessed.
    :param in_dem_file: the DEM file which is to be used for processing. Either:
                        * GDAL readable file which will be subset/reprojected and converted to gamma binary
                        * Gamma binary DEM file which corresponds to the gamma dem par file.
    :param demparfile: The Gamma dem par file. If the file exists and use_dem_file is True then the in_dem_file
                       will be used. If check_in_dem_filename is True then the file name for the in_dem_file will be
                       derived from the demparfile name by removing extension and replacing with extension .dem.
    :param outbasename: The base name for the output file. All output files will be outbasename + process_stage.ext
    :param out_dir: A directory where final geotiff files generated by the processing chain will be stored.
    :param tmp_dir: A directory where temporary output files generated by the processing chain will be stored.
    :param out_res_x: A float with the output ARD resolution in the x axis.
    :param out_res_y: A float within the output ARD resolution in the y axis.
    :param out_proj_epsg: An int for the EPSG code for the output projection. If None then UTM will be used with the
                          zone automatically derived from the centre point of the scene.
    :param use_dem_file: Boolean. If True the DEM file is expected to be available within Gamma format.
    :param check_in_dem_filename: If True then the DEM file name will be derived from the demparfile (by removing
                                  extension and adding .dem).
    :param no_dem_check: If True the DEM is not checked for no data value and minimum value for Gamma compatibility
    :param dem_resample_method: specify the gdal resampling method for the DEM processing (if DEM being subsetted
                                and reprojected). Default: gdal.GRA_CubicSpline

    """
    logger.info("Starting processing of {}.".format(sen1img))

    # -------------------------- STEP 0: Debug record function inputs -------------------------- #
    logger.debug("Annotation XML file: {}.".format(ann_sen1_xml))
    logger.debug("Calibration XML file: {}.".format(cal_sen1_xml))
    logger.debug("Noise XML file: {}.".format(nos_sen1_xml))
    logger.debug("Input DEM file: {}.".format(in_dem_file))
    logger.debug("DEM Par file: {}.".format(demparfile))
    logger.debug("Output file basename: {}.".format(outbasename))
    logger.debug("Output directory: {}.".format(out_dir))
    logger.debug("Temporary directory: {}.".format(tmp_dir))
    logger.debug("Output resolution X axis: {}.".format(out_res_x))
    logger.debug("Output resolution Y axis: {}.".format(out_res_y))
    logger.debug("Output EPSG projection: {}.".format(out_proj_epsg))
    logger.debug("use_dem_file: {}.".format(use_dem_file))
    logger.debug("check_in_dem_filename: {}.".format(check_in_dem_filename))
    logger.debug("dem_resample_method: {}.".format(dem_resample_method))
    # -------------------------- STEP 0: End -------------------------- #

    # -------------------------- STEP 1: Import DEM -------------------------- #
    logger.info("Import DEM")
    if os.path.exists(demparfile) and use_dem_file:
        logger.debug("Expecting DEM to already be in Gamma format...")
        if check_in_dem_filename:
            dem_test_file = os.path.splitext(demparfile)[0] + '.dem'
            logger.debug("Expecting DEM to have file name: '{}'".format(dem_test_file))
            if os.path.exists(dem_test_file):
                demfile = dem_test_file
        else:
            logger.debug("Expecting DEM to have given file name: '{}'".format(in_dem_file))
            demfile = in_dem_file
    else:
        logger.debug("Expecting DEM to be in GDAL readable format: '{}'".format(in_dem_file))
        dem_reproj_file = os.path.splitext(demparfile)[0] + '.dem'
        logger.debug("Output DEM will have file name: '{}'".format(dem_reproj_file))

        sen1_bbox_wgs84 = get_sen1_latlong_bbox(ann_sen1_xml)
        logger.debug("WGS84 BBOX of input Sentinel 1 scene: [{}, {}, {}, {}]".format(sen1_bbox_wgs84[0],
                                                                                sen1_bbox_wgs84[1],
                                                                                sen1_bbox_wgs84[2],
                                                                                sen1_bbox_wgs84[3]))

        if out_proj_epsg is None:
            logger.debug("Finding UTM zone for scene from centre point")
            centre_lat = sen1_bbox_wgs84[2] + ((sen1_bbox_wgs84[3] - sen1_bbox_wgs84[2]) / 2)
            centre_lon = sen1_bbox_wgs84[0] + ((sen1_bbox_wgs84[1] - sen1_bbox_wgs84[0]) / 2)
            logger.debug("WGS84 Sentinel 1 scene centre point: [{}, {}]".format(centre_lat, centre_lon))
            utm_zone = sen1_ard_gamma.sen1_ard_utils.latlon_to_zone_number(centre_lat, centre_lon)
            logger.debug("UTM Zone for Scene: {}".format(utm_zone))
            out_proj_epsg = sen1_ard_gamma.sen1_ard_utils.epsg_for_UTM(utm_zone, 'N')

        logger.debug("EPSG Code for output: {}".format(out_proj_epsg))

        sen1_ard_gamma.gamma_dem_prep.subset_reproj_utm_dem_file(in_dem_file, sen1_bbox_wgs84, out_res_x, out_res_y,
                                                                 out_proj_epsg, dem_reproj_file, demparfile, tmp_dir,
                                                                 no_dem_check, eResampleAlg=dem_resample_method)

        demfile = dem_reproj_file
    logger.debug("Finished Importing DEM.")
    # -------------------------- STEP 1: End -------------------------- #

    # -------------------------- STEP 2: Import to Gamma -------------------------- #
    logger.info("Import Sentinel-1 Scene to Gamma.")
    sen1_img_import = os.path.join(tmp_dir, outbasename + '_import')
    sen1_par_import = os.path.join(tmp_dir, outbasename + '_import.par')
    sen1_nesz_import = os.path.join(tmp_dir, outbasename + '_import_nesz')

    cmd = 'par_S1_GRD {0} {1} {2} {3} {4} {5} - - 1 - {6}'.format(sen1img, ann_sen1_xml, cal_sen1_xml, nos_sen1_xml,
                                                                  sen1_par_import, sen1_img_import, sen1_nesz_import)

    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug("Finished importing Sentinel-1 Scene to Gamma.")
    # -------------------------- STEP 2: End -------------------------- #

    # ------ STEP 3: Read resolution and inc. angle info from image and dem parameter file ------------ #
    pardict = sen1_ard_gamma.sen1_ard_utils.gamma_par_dict(sen1_par_import)
    mli_rgres = float(pardict['range_pixel_spacing'][0])
    mli_azres = float(pardict['azimuth_pixel_spacing'][0])
    center_lat = float(pardict['center_latitude'][0])
    inc = float(pardict['incidence_angle'][0])

    demdict = sen1_ard_gamma.sen1_ard_utils.gamma_par_dict(demparfile)
    proj = demdict['DEM_projection'][0]
    if proj == 'EQA':
        dem_xres = float(demdict['post_lon'][0])
        dem_yres = abs(float(demdict['post_lat'][0]))
    else:
        dem_xres = float(demdict['post_east'][0])
        dem_yres = abs(float(demdict['post_north'][0]))
    # -------------------------- STEP 3: End -------------------------- #

    # -------------------------- STEP 4: Calculate DEM oversampling factors -------------------------- #
    ovr_lon = abs(dem_xres / out_res_y)
    ovr_lat = abs(dem_yres / out_res_x)
    # -------------------------- STEP 4: End -------------------------- #

    # -------------------------- STEP 5: Calculate multilooking factors -------------------------- #
    if proj == 'EQA':
        deg_in_m = 111319 * math.cos(center_lat * 3.141 / 180)
        geo_xres_m = deg_in_m * out_res_y
        geo_yres_m = 110946 * abs(out_res_x)
        ml_rg = int(math.ceil(geo_xres_m / mli_rgres))
        ml_az = int(math.ceil(geo_yres_m / mli_azres))
        if ((ml_rg * mli_rgres) / geo_xres_m) > 1.25:
            ml_rg = int(math.floor(geo_xres_m / mli_rgres))
        if ((ml_az * mli_azres) / geo_yres_m) > 1.25:
            ml_az = int(math.floor(geo_yres_m / mli_azres))
    else:
        ml_rg = int(math.ceil(out_res_y / mli_rgres))
        ml_az = int(math.ceil(out_res_x / mli_azres))
        if ((ml_rg * mli_rgres) / out_res_y) > 1.25:
            ml_rg = int(math.floor(out_res_y / mli_rgres))
        if ((ml_az * mli_azres) / out_res_x) > 1.25:
            ml_az = int(math.floor(out_res_x / mli_azres))

            # In case calculation results in ml-factors of 0
    if (ml_rg < 1):
        ml_rg = 1
    if (ml_az < 1):
        ml_az = 1
    # -------------------------- STEP 5: End -------------------------- #

    # -------------------------- STEP 6: Multilooking -------------------------- #
    logger.info("Multilooking: {0} x {1}".format(ml_rg, ml_az))

    sen1_img_mli = os.path.join(tmp_dir, outbasename + '_import_mli')
    sen1_par_mli = os.path.join(tmp_dir, outbasename + '_import_mli.par')

    cmd = 'multi_look_MLI {0} {1} {2} {3} {4} {5}'.format(sen1_img_import, sen1_par_import, sen1_img_mli, sen1_par_mli,
                                                          ml_rg, ml_az)

    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    logger.debug("Finished Multilooking: {0} x {1}".format(ml_rg, ml_az))
    # -------------------------- STEP 6: End -------------------------- #

    # -------------------------- STEP 7: Read width and lines from mlipar -------------------------- #
    mlidict = sen1_ard_gamma.sen1_ard_utils.gamma_par_dict(sen1_par_mli)
    mli_width = int(mlidict['range_samples'][0])
    mli_lines = int(mlidict['azimuth_lines'][0])
    # -------------------------- STEP 7: End -------------------------- #

    # -------------------------- STEP 8: Subtract noise equivalent sigma nought from MLI -------------------------- #
    logger.info("Multilook S1 GRD noise image...")
    sen1_nesz_mli = os.path.join(tmp_dir, outbasename + '_import_nesz_mli')
    sen1_nesz_par_mli = os.path.join(tmp_dir, outbasename + '_import_nesz_mli.par')
    cmd = 'multi_look_MLI {0} {1} {2} {3} {4} {5}'.format(sen1_nesz_import, sen1_par_import, sen1_nesz_mli,
                                                          sen1_nesz_par_mli, ml_rg, ml_az)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug("Finished multilook S1 GRD noise image...")

    logger.info('Subtract NESZ from MLI')
    sen1_img_mli_cor = os.path.join(tmp_dir, outbasename + '_import_mli_cor')

    cmd = 'float_math {0} {1} {2} {3} 1 - - - - - 1'.format(sen1_img_mli, sen1_nesz_mli, sen1_img_mli_cor, mli_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.info('Finished Subtract NESZ from MLI')

    logger.info("Remove pwr values below -40 dB")
    sen1_img_mli_cor2 = os.path.join(tmp_dir, outbasename + '_import_mli_cor2')
    cmd = 'replace_values {0} 0.0001 0.0001 {1} {2} 2 2'.format(sen1_img_mli_cor, sen1_img_mli_cor2, mli_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_mli = sen1_img_mli_cor2
    logger.debug("Finished remove pwr values below -40 dB")
    # -------------------------- STEP 8: End -------------------------- #

    # -------------------------- STEP 9: Calibration -------------------------- #
    logger.info("Running calibration")
    sen1_img_cmli = os.path.join(tmp_dir, outbasename + '_import_cmli')
    sen1_img_pix_ell = os.path.join(tmp_dir, outbasename + '_import_pix_ell')
    cmd = 'radcal_MLI {0} {1} - {2} - 0 0 1 0 - {3}'.format(sen1_img_mli, sen1_par_mli, sen1_img_cmli, sen1_img_pix_ell)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_cmli = sen1_img_mli
    logger.debug("Finished calibration")
    # -------------------------- STEP 9: End -------------------------- #

    # -------------------------- STEP 10: Compute Geocoding Lookup Table -------------------------- #
    logger.info('Compute Geocoding Lookup Table')

    sen1_img_inc = os.path.join(tmp_dir, outbasename + '_import_inc')
    sen1_img_pix = os.path.join(tmp_dir, outbasename + '_import_pix')
    sen1_img_lsmap = os.path.join(tmp_dir, outbasename + '_import_lsmap')
    sen1_img_geo2rdc = os.path.join(tmp_dir, outbasename + '_import_geo2rdc')
    sen1_img_simsar = os.path.join(tmp_dir, outbasename + '_import_simsar')
    sen1_img_gcdem = os.path.join(tmp_dir, outbasename + '_import_gcdem')
    sen1_par_gcdem = os.path.join(tmp_dir, outbasename + '_import_gcdem.par')

    cmd = 'gc_map {0} - {1} {2} {3} {4} {5} {6} {7} {8} - - {9} - {10} {11} - 2'.format(sen1_par_mli, demparfile,
                                                                                demfile, sen1_par_gcdem,
                                                                                sen1_img_gcdem,
                                                                                sen1_img_geo2rdc, ovr_lon,
                                                                                ovr_lat, sen1_img_simsar,
                                                                                sen1_img_inc, sen1_img_pix,
                                                                                sen1_img_lsmap)

    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    logger.debug('Finished Computing Geocoding Lookup Table')
    # -------------------------- STEP 10: End -------------------------- #

    # -------------------------- STEP 11: Read width and lines from sen1_par_gcdem -------------------------- #
    gcdemdict = sen1_ard_gamma.sen1_ard_utils.gamma_par_dict(sen1_par_gcdem)
    dem_width = int(gcdemdict['width'][0])
    # -------------------------- STEP 11: End -------------------------- #

    # -------------------------- STEP 12: Estimate pixel scattering area based on DEM -------------------------- #
    logger.info('Estimate pixel scattering area based on DEM')

    sen1_img_pixdem = os.path.join(tmp_dir, outbasename + '_import_pix_dem')

    cmd = 'pixel_area {0} {1} {2} {3} {4} {5} {6}'.format(sen1_par_mli, sen1_par_gcdem, sen1_img_gcdem,
                                                          sen1_img_geo2rdc, sen1_img_lsmap, sen1_img_inc,
                                                          sen1_img_pixdem)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    logger.debug('Finished estimating pixel scattering area based on DEM')
    # -------------------------- STEP 12: End -------------------------- #

    # -------------------------- STEP 13: Improve Geocoding Lookup Table -------------------------- #
    logger.info('Check Accuracy of Geocoding Lookup Table')

    sen1_par_diff = os.path.join(tmp_dir, outbasename + '_import_diff.par')
    cmd = 'create_diff_par {0} - {1} 1 0'.format(sen1_par_mli, sen1_par_diff)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_offs = os.path.join(tmp_dir, outbasename + '_import_offs')
    sen1_img_ccp = os.path.join(tmp_dir, outbasename + '_import_ccp')
    sen1_img_offsets = os.path.join(tmp_dir, outbasename + '_import_offsets')
    cmd = 'offset_pwrm {0} {1} {2} {3} {4} 128 128 {5} - 64 64 0.2'.format(sen1_img_cmli, sen1_img_pixdem,
                                                                           sen1_par_diff, sen1_img_offs, sen1_img_ccp,
                                                                           sen1_img_offsets)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_coffs = os.path.join(tmp_dir, outbasename + '_import_coffs')
    sen1_img_coffsets = os.path.join(tmp_dir, outbasename + '_import_coffsets')
    cmd = 'offset_fitm {0} {1} {2} {3} {4} 0.2 - 0'.format(sen1_img_offs, sen1_img_ccp, sen1_par_diff, sen1_img_coffs,
                                                           sen1_img_coffsets)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    logger.debug('Finished Checking Accuracy of Geocoding Lookup Table')
    # -------------------------- STEP 13: End -------------------------- #

    # -------------------------- STEP 14: Perform Fine Registration -------------------------- #
    logger.info('Improve Geocoding Lookup Table')

    sen1_img_geo2rdc_fine = os.path.join(tmp_dir, outbasename + '_import_geo2rdc_fine')
    cmd = 'gc_map_fine {0} {1} {2} {3} 0'.format(sen1_img_geo2rdc, dem_width, sen1_par_diff, sen1_img_geo2rdc_fine)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.info('Finished improving Geocoding Lookup Table')
    sen1_img_geo2rdc = sen1_img_geo2rdc_fine

    logger.info('Update Pixel Area map')
    cmd = 'pixel_area {0} {1} {2} {3} {4} {5} {6}'.format(sen1_par_mli, sen1_par_gcdem, sen1_img_gcdem,
                                                          sen1_img_geo2rdc, sen1_img_lsmap, sen1_img_inc,
                                                          sen1_img_pixdem)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug('Finished updating Pixel Area map')
    # -------------------------- STEP 14: End -------------------------- #

    # -------------------------- STEP 15: Pixel Area Correction -------------------------- #
    logger.info('Pixel Area Correction')

    cmd = 'float_math {0} {1} {2} {3} 3 - - 1 1 - '.format(sen1_img_pix_ell, sen1_img_pixdem, sen1_img_pix, mli_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_s0 = os.path.join(tmp_dir, outbasename + '_import_s0')
    cmd = 'float_math {0} {1} {2} {3}  2 - - 1 1 - '.format(sen1_img_cmli, sen1_img_pix, sen1_img_s0, mli_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug('Finished Pixel Area Correction')
    # -------------------------- STEP 15: End -------------------------- #

    # -------------------------- STEP 16: Geocode and pixel area correction for GRDs -------------------------- #
    logger.info('Geocode to projection: {}'.format(proj))
    sen1_img_s0geo = os.path.join(tmp_dir, outbasename + '_import_s0geo')
    cmd = 'geocode_back {0} {1} {2} {3} {4} - 7 0'.format(sen1_img_s0, mli_width, sen1_img_geo2rdc, sen1_img_s0geo,
                                                          dem_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_pixgeo = os.path.join(tmp_dir, outbasename + '_import_pixgeo')
    cmd = 'geocode_back {0} {1} {2} {3} {4} - 1 0 '.format(sen1_img_pix, mli_width, sen1_img_geo2rdc, sen1_img_pixgeo,
                                                           dem_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    sen1_img_pix = sen1_img_pixgeo
    logger.debug('Finished geocode to projection: {}'.format(proj))
    # -------------------------- STEP 16: End -------------------------- #

    # -------------------------- STEP 17: Sigma to Gamma conversion -------------------------- #
    logger.info('Convert Sigma to gamma.')
    sen1_img_g0geo = os.path.join(tmp_dir, outbasename + '_import_g0geo')
    cmd = 'sigma2gamma {0} {1} {2} {3} '.format(sen1_img_s0geo, sen1_img_inc, sen1_img_g0geo, dem_width)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug('Finished converting Sigma to gamma.')

    # -------------------------- STEP 17: End -------------------------- #

    # -------------------------- STEP 18: Export to geotiff -------------------------- #
    logger.info('Convert file outputs to GeoTiff')

    sen1_img_pwrout = os.path.join(out_dir, outbasename + '_final_pwr.tif')
    sen1_img_incout = os.path.join(out_dir, outbasename + '_final_inc.tif')
    sen1_img_lsout = os.path.join(out_dir, outbasename + '_final_lsmap.tif')
    sen1_img_pixout = os.path.join(out_dir, outbasename + '_final_pix.tif')

    cmd = 'data2geotiff {0} {1} 2 {2}'.format(sen1_par_gcdem, sen1_img_g0geo, sen1_img_pwrout)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    cmd = 'data2geotiff {0} {1} 2 {2}'.format(sen1_par_gcdem, sen1_img_inc, sen1_img_incout)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    cmd = 'data2geotiff {0} {1} 2 {2}'.format(sen1_par_gcdem, sen1_img_pix, sen1_img_pixout)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e

    cmd = 'data2geotiff {0} {1} 2 {2}'.format(sen1_par_gcdem, sen1_img_lsmap, sen1_img_lsout)
    try:
        logger.debug("Running following command using subprocess '{}'".format(cmd))
        cmd = sen1_ard_gamma.sen1_ard_utils.preappend_cmd(cmd)
        subprocess.call(cmd, shell=True)
    except OSError as e:
        logger.error("Command failed: '{}'".format(cmd))
        raise e
    logger.debug('Finished converting file outputs to GeoTiff')
    # -------------------------- STEP 18: End -------------------------- #

    logger.info("Completed processing of {}.".format(sen1img))

def create_pol_stacked_products(out_scns_dict, out_base_name, tmp_dir, out_dir, gdal_format='GTIFF', out_int_imgs=False):
    """
    A function to create the final stacked ARD image products.

    :param out_scns_dict: dictionary with polarisation at the first level and second level of (inc, lsmap, pix, pwr)
    :param out_base_name: output image file base name
    :param tmp_dir: tmp directory to sort tempory images to
    :param out_dir: output directory to export final images to
    :param gdal_format: the output GDAL image format (Default: GTIFF)

    """
    logger.info("Starting to stack images to create final products")
    out_img_list = dict()
    img_ext = sen1_ard_gamma.sen1_ard_utils.get_file_extension(gdal_format)
    scn_keys = list(out_scns_dict.keys())

    dB_pwr_img = ""
    if len(scn_keys) == 2:
        if ('vv' in scn_keys) and ('vh' in scn_keys):
            # ------------ Process inc Product ------------- #
            out_img = os.path.join(out_dir, "{}_inc{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain.
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_scns_dict['vv']['inc'], out_img, gdal_format,
                                                                10000, numpy.uint16, 0, 0)
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_scns_dict['vv']['inc'],  out_img, gdal_format)
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['inc'])
            out_img_list['inc'] = out_img
            # ------------ Processed inc Product ------------- #

            # ------------ Process pix Product ------------- #
            # Create tmp VRT stacked image
            out_vrt_tmp = os.path.join(tmp_dir, "{}_pix_tmp.vrt".format(out_base_name))
            pix_imgs = [out_scns_dict['vv']['pix'], out_scns_dict['vh']['pix']]
            sen1_ard_gamma.sen1_ard_utils.gdal_stack_images_vrt(pix_imgs, out_vrt_tmp)

            # Generate final output image
            out_img = os.path.join(out_dir, "{}_pix{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain.
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_vrt_tmp, out_img, gdal_format,
                                                                10000, numpy.uint16, 0, 0)
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_vrt_tmp, out_img, gdal_format)
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['vv', 'vh'])
            out_img_list['pix'] = out_img
            # ------------ Processed pix Product ------------- #

            # ------------ Process pwr Product ------------- #
            out_ratio_tmp_file = os.path.join(tmp_dir, "{}_pwr_ratio_tmp{}".format(out_base_name, img_ext))
            sen1_ard_gamma.sen1_ard_utils.calc_ratio_img(out_scns_dict['vv']['pwr'], out_scns_dict['vh']['pwr'],
                                                         out_ratio_tmp_file, gdal_format)

            # Create tmp VRT stacked power image.
            out_vrt_tmp = os.path.join(tmp_dir, "{}_pwr_tmp.vrt".format(out_base_name))
            pwr_imgs = [out_scns_dict['vv']['pwr'], out_scns_dict['vh']['pwr'], out_ratio_tmp_file]
            sen1_ard_gamma.sen1_ard_utils.gdal_stack_images_vrt(pwr_imgs, out_vrt_tmp)

            # Generate final output image
            out_img = os.path.join(out_dir, "{}_pwr{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_vrt_tmp, out_img, gdal_format,
                                                                100000, numpy.uint32, 0, 0)
                dB_pwr_img = out_vrt_tmp
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_vrt_tmp, out_img, gdal_format)
                dB_pwr_img = out_img
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['vv', 'vh', 'vv/vh'])
            out_img_list['pwr'] = out_img
            # ------------ Processed pwr Product ------------- #
        elif ('hh' in scn_keys) and ('hv' in scn_keys):
            # ------------ Process inc Product ------------- #
            out_img = os.path.join(out_dir, "{}_inc{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain.
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_scns_dict['hh']['inc'], out_img, gdal_format,
                                                                10000, numpy.uint16, 0, 0)
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_scns_dict['hh']['inc'],  out_img, gdal_format)
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['inc'])
            out_img_list['inc'] = out_img
            # ------------ Processed inc Product ------------- #

            # ------------ Process pix Product ------------- #
            # Create tmp VRT stacked image
            out_vrt_tmp = os.path.join(tmp_dir, "{}_pix_tmp.vrt".format(out_base_name))
            pix_imgs = [out_scns_dict['hh']['pix'], out_scns_dict['hv']['pix']]
            sen1_ard_gamma.sen1_ard_utils.gdal_stack_images_vrt(pix_imgs, out_vrt_tmp)

            # Generate final output image
            out_img = os.path.join(out_dir, "{}_pix{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain.
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_vrt_tmp, out_img, gdal_format,
                                                                10000, numpy.uint16, 0, 0)
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_vrt_tmp, out_img, gdal_format)
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['hh', 'hv'])
            out_img_list['pix'] = out_img
            # ------------ Processed pix Product ------------- #

            # ------------ Process pwr Product ------------- #
            out_ratio_tmp_file = os.path.join(tmp_dir, "{}_pwr_ratio_tmp{}".format(out_base_name, img_ext))
            sen1_ard_gamma.sen1_ard_utils.calc_ratio_img(out_scns_dict['hh']['pwr'], out_scns_dict['hv']['pwr'],
                                                         out_ratio_tmp_file, gdal_format)

            # Create tmp VRT stacked power image.
            out_vrt_tmp = os.path.join(tmp_dir, "{}_pwr_tmp.vrt".format(out_base_name))
            pwr_imgs = [out_scns_dict['hh']['pwr'], out_scns_dict['hv']['pwr'], out_ratio_tmp_file]
            sen1_ard_gamma.sen1_ard_utils.gdal_stack_images_vrt(pwr_imgs, out_vrt_tmp)

            # Generate final output image
            out_img = os.path.join(out_dir, "{}_pwr{}".format(out_base_name, img_ext))
            if out_int_imgs:
                # Multiply by Gain
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_vrt_tmp, out_img, gdal_format,
                                                                100000, numpy.uint32, 0, 0)
                dB_pwr_img = out_vrt_tmp
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_vrt_tmp, out_img, gdal_format)
                dB_pwr_img = out_img
            sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, ['hh', 'hv', 'hh/hv'])
            out_img_list['pwr'] = out_img
            # ------------ Processed pwr Product ------------- #
        else:
            raise Exception("Do not recognise the polarisations provided")
    elif len(scn_keys) == 1:
        pol = scn_keys[0]
        for prod in ['inc', 'pix', 'pwr']:
            out_img = os.path.join(out_dir, "{}_{}{}".format(out_base_name, prod, img_ext))
            if out_int_imgs:
                # Multiply by Gain
                gain_val = 10000
                np_dtype = numpy.uint16
                if (prod == 'pwr'):
                    gain_val = 100000
                    np_dtype = numpy.uint32
                sen1_ard_gamma.sen1_ard_utils.apply_gain_to_img(out_scns_dict[pol][prod], out_img, gdal_format,
                                                                gain_val, np_dtype, 0, 0)
            else:
                # GDAL Translate to output file
                sen1_ard_gamma.sen1_ard_utils.gdal_translate(out_scns_dict[pol][prod], out_img, gdal_format)
            if prod in ['inc', 'lsmap']:
                sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, [prod])
            else:
                sen1_ard_gamma.sen1_ard_utils.set_band_names(out_img, [pol])
            out_img_list[prod] = out_img
        dB_pwr_img = out_scns_dict[pol]['pwr']
    else:
        raise Exception("Only know how to stack images with 2 bands or copy single band images.")
    if 'pwr' in out_img_list:
        logger.debug("Calculate dB image from power.")
        out_dB_img = os.path.join(out_dir, "{}_dB{}".format(out_base_name, img_ext))
        sen1_ard_gamma.sen1_ard_utils.convert_to_dB(dB_pwr_img, out_dB_img, gdal_format, out_int_imgs)
        out_img_list['dB'] = out_dB_img

        out_vmsk_img = os.path.join(out_dir, "{}_vmsk{}".format(out_base_name, img_ext))
        sen1_ard_gamma.sen1_ard_utils.calc_valid_msk(dB_pwr_img, out_vmsk_img, gdal_format, 0.0)
        out_img_list['vmsk'] = out_vmsk_img

    return out_img_list

def run_sen1_grd_ard_analysis(input_safe_file, output_dir, tmp_dir, dem_img_file, out_img_res, out_proj_epsg,
                              polarisations, gdal_format, calc_no_stats, keep_files, no_dem_check=False,
                              out_int_imgs=False):
    """
    High level function which runs the analysis of a Sentinel-1 GRD scene to an ARD product.

    :param input_safe_file:
    :param output_dir:
    :param tmp_dir:
    :param dem_img_file:
    :param out_img_res:
    :param out_proj_epsg:
    :param polarisations:
    :param gdal_format:
    :param calc_no_stats:
    :param keep_files:
    :param no_dem_check: If True the DEM is not checked for no data value and minimum value for Gamma compatibility
    :param out_int_imgs: Produce intergised output images with gains applied.
    """
    scn_metadata_info = sen1_ard_gamma.sen1_ard_utils.retrieve_sentinel1_metadata(input_safe_file)

    scn_safe_files = sen1_ard_gamma.sen1_ard_utils.find_sen1_ard_files(input_safe_file)

    scn_basename = sen1_ard_gamma.sen1_ard_utils.create_sentinel1_basename(scn_metadata_info)
    logger.info("Basename for scene: {}".format(scn_basename))

    if polarisations is None:
        polarisations = scn_metadata_info['product_polarisations']
    else:
        for pol in polarisations:
            if pol not in scn_metadata_info['product_polarisations']:
                raise Exception("Polarisation {} is not within the scene provided.".format(pol))

    c_uid = sen1_ard_gamma.sen1_ard_utils.uidGenerator()
    c_tmp_dir = os.path.join(tmp_dir, '{}_tmp_{}'.format(scn_basename, c_uid))
    c_tmp_dir_created = False
    if not os.path.exists(c_tmp_dir):
        os.mkdir(c_tmp_dir)
        c_tmp_dir_created = True

    c_out_dir = os.path.join(tmp_dir, '{}_{}'.format(scn_basename, c_uid))
    c_out_dir_created = False
    if not os.path.exists(c_out_dir):
        os.mkdir(c_out_dir)
        c_out_dir_created = True

    demparfile = os.path.join(c_tmp_dir, scn_basename + '_gamma_scn_dem.dem_par')

    out_scns = dict()
    first = True
    for pol in polarisations:
        pol_lower = pol.lower()
        c_scn_basename = scn_basename + '_' + pol_lower
        logger.info("Processing {} Polarisation.".format(pol))

        sen1_ard_gamma.sen1_grd_ard_tools.exe_gamma_grd_ard_processing(scn_safe_files['measure_' + pol_lower],
                                                                       scn_safe_files['annotation_' + pol_lower],
                                                                       scn_safe_files['calibration_' + pol_lower],
                                                                       scn_safe_files['noise_' + pol_lower],
                                                                       dem_img_file, demparfile, c_scn_basename,
                                                                       c_out_dir, c_tmp_dir, out_img_res, -out_img_res,
                                                                       out_proj_epsg, use_dem_file=(not first),
                                                                       check_in_dem_filename=True,
                                                                       no_dem_check=no_dem_check,
                                                                       dem_resample_method=gdal.GRA_CubicSpline)

        out_scns[pol_lower] = dict()
        out_files = glob.glob(os.path.join(c_out_dir, "{}*.tif".format(c_scn_basename)))
        for img_file in out_files:
            img_file_name = os.path.basename(img_file)
            if 'inc' in img_file_name:
                out_scns[pol_lower]['inc'] = img_file
            elif 'lsmap' in img_file_name:
                out_scns[pol_lower]['lsmap'] = img_file
            elif 'pix' in img_file_name:
                out_scns[pol_lower]['pix'] = img_file
            elif 'pwr' in img_file_name:
                out_scns[pol_lower]['pwr'] = img_file
        first = False

    logger.info("Creating final image outputs")
    # Create final output images with bands stacked where appropriate and ratio image (vv/vh) calculated.
    fnl_out_imgs = sen1_ard_gamma.sen1_grd_ard_tools.create_pol_stacked_products(out_scns, scn_basename, c_out_dir,
                                                                                output_dir, gdal_format, out_int_imgs)

    if not calc_no_stats:
        logger.info("Calculating image statistics and pyramids for final outputs")
        if 'vmsk' in fnl_out_imgs:
            sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['vmsk'], no_data_val=0.0)
        if 'inc' in fnl_out_imgs:
            sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['inc'], no_data_val=0.0)
        if 'pix' in fnl_out_imgs:
            sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['pix'], no_data_val=0.0)
        if 'pwr' in fnl_out_imgs:
            sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['pwr'], no_data_val=0.0)
        if 'dB' in fnl_out_imgs:
            if out_int_imgs:
                sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['dB'], no_data_val=32767)
            else:
                sen1_ard_gamma.calc_img_stats.run_calc_img_stats_pyramids(fnl_out_imgs['dB'], no_data_val=999)

    if not keep_files:
        if c_tmp_dir_created:
            shutil.rmtree(c_tmp_dir)
        if c_out_dir_created:
            shutil.rmtree(c_out_dir)

    logger.info("Completed all data processing stages; outputs in {}".format(output_dir))


text:
import os.path as osp
import numpy as np
import sys

import zarr
from dl_ext.vision_ext.datasets.kitti.load import load_calib
from tqdm import trange

sys.path.append('deepv2d')

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

import cv2
import os
import time
import argparse

# from deepv2d import eval_utils
from deepv2d.core import config
from deepv2d.deepv2d import DeepV2D
from deepv2d.data_stream.kitti import KittiRaw

# def process_for_evaluation(depth, scale, crop):
#     """ During training ground truth depths are scaled and cropped, we need to
#         undo this for evaluation """
#     depth = (1.0 / scale) * np.pad(depth, [[crop, 0], [0, 0]], 'mean')
#     return depth

def img_to_rect(u, v, depth_rect, K):
    fu, fv, cu, cv = K
    x = ((u - cu) * depth_rect) / fu
    y = ((v - cv) * depth_rect) / fv
    pts_rect = np.concatenate((x.reshape(-1, 1), y.reshape(-1, 1), depth_rect.reshape(-1, 1)), axis=1)
    return pts_rect

def depth_to_rect(depth_map, intrinsic):
    x_range = np.arange(0, depth_map.shape[1])
    y_range = np.arange(0, depth_map.shape[0])
    x_idxs, y_idxs = np.meshgrid(x_range, y_range)
    x_idxs, y_idxs = x_idxs.reshape(-1), y_idxs.reshape(-1)
    depth = depth_map[y_idxs, x_idxs]
    pts_rect = img_to_rect(x_idxs, y_idxs, depth, intrinsic)
    return pts_rect

def make_predictions(args):
    """ Run inference over the test images """

    np.random.seed(1234)
    cfg = config.cfg_from_file(args.cfg)

    db = KittiRaw(args.dataset_dir, split='object')
    scale = db.args['scale']
    crop = db.args['crop']

    os.makedirs(args.output_dir, exist_ok=True)

    deepv2d = DeepV2D(cfg, args.model, use_fcrn=False, mode='keyframe')
    tfconfig = tf.ConfigProto()
    tfconfig.gpu_options.allow_growth = True

    with tf.Session(config=tfconfig) as sess:
        deepv2d.set_session(sess)

        predictions = []
        # it = db.test_set_iterator()
        for i in trange(7481):
            # for i, data_blob in enumerate():
            try:
                data_blob = db.test_set_iterator(i)
            except FileNotFoundError:
                # np.empty((0, 4)).tofile(osp.join(args.output_dir, '%06d.bin' % i))
                zarr.save(osp.join(args.output_dir, '%06d.zarr' % i), np.zeros(()))
                continue
            images, intrinsics, poses = data_blob['images'], data_blob['intrinsics'], data_blob['poses']
            if not args.input_pose:
                poses = None
                iters = args.n_iters
            else:
                iters = 1
            depth_predictions, _ = deepv2d(images, intrinsics,
                                           iters=iters, camposes=poses)
            keyframe_depth = depth_predictions[0]
            keyframe_depth = keyframe_depth / scale

            zarr.save(osp.join(args.output_dir, '%06d.zarr' % i), keyframe_depth)

            # keyframe_image = images[0]
            #
            # pts = depth_to_rect(keyframe_depth, intrinsics)
            # calib = load_calib('/home/linghao/Datasets/kitti', 'training', i)
            # lidar = calib.rect_to_lidar(pts)
            # lidar = np.concatenate((lidar, np.ones((lidar.shape[0], 1))), axis=1)
            # lidar.tofile(osp.join(args.output_dir, '%06d.bin' % i))

        return predictions

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', default='cfgs/kitti.yaml', help='config file used to train the model')
    parser.add_argument('--model', default='models/kitti.ckpt', help='path to model checkpoint')
    parser.add_argument('--dataset_dir', default='/home/linghao/Datasets/kitti/kitti_raw',
                        help='config file used to train the model')
    parser.add_argument('--viz', action="store_true", help='display depth maps during inference')
    parser.add_argument('--n_iters', type=int, default=5, help='number of video frames to use for reconstruction')
    parser.add_argument('--output_dir', type=str, default='kitti_object_depth')
    parser.add_argument('--input_pose', default=True)
    args = parser.parse_args()

    # run inference on the test images
    predictions = make_predictions(args)

ext: py

reason: import numpy


text:
"""Extract column data using grok patterns"""

import datatable as dt
import numpy as np
from h2oaicore.transformer_utils import CustomTransformer

#
# Extract data from a composite column record using grok mapping a la logstash
#
# Example pattern:
_PATTERN = '%{TIMESTAMP_ISO8601:ts} %{IPV4:ip} %{NUMBER:status:int}: %{LOGLEVEL:level} %{WORD:class} %{DATA}'
#
# Column name to transform using the pattern above
_COLUMN_TO_PARSE = 'syslog'
#
#
# These new columns will be added based on the pattern above:
# syslog_ts, syslog_ip, syslog_status, syslog_level, syslog_class
#
# CSV example:
#
# column1,syslog,some_result
# hello,"2011-09-01 11:00:33.4444 10.10.10.1 200: INFO class1 Some data here 1",0
# hello,"2011-09-02 12:00:33.4444 10.10.10.1 200: WARN class2 Some data here 2",0
#

class TextGrokParser(CustomTransformer):

    _modules_needed_by_name = ["pygrok==1.0.0"]

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        from pygrok import Grok
        self.grok = Grok(_PATTERN)
        self.columns = list(self.grok.regex_obj.groupindex.keys())
        self.column_to_parse = _COLUMN_TO_PARSE

    @staticmethod
    def get_default_properties():
        return dict(col_type="text", min_cols=1, max_cols=1, relative_importance=1)

    @staticmethod
    def do_acceptance_test():
        return False

    # noinspection PyBroadException
    def parse_text(self, text):
        try:
            x = self.grok.match(text)
            return tuple(x[k] for k in self.columns)
        except Exception:
            return tuple(None for _ in self.columns)

    def fit_transform(self, X: dt.Frame, y: np.array = None):
        return self.transform(X)

    def transform(self, X: dt.Frame):
        col_name = X.names[0]
        if col_name == self.column_to_parse:
            Y = X[col_name].to_list()[0]
            Z = dt.Frame([self.parse_text(x) for x in Y], names=[f"{col_name}_{s}" for s in self.columns])
            return Z
        else:
            return X.to_pandas().iloc[:, None]


text:
import os
import pathlib
import yaml
import copy
import numpy as np

from aiida.engine import WorkChain, ToContext, ExitCode
from aiida.orm import Int, Bool, Code, Dict, List
from aiida.orm import SinglefileData, StructureData
from aiida.plugins import WorkflowFactory
from aiida_nanotech_empa.workflows.cp2k.cp2k_utils import get_kinds_section, determine_kinds, dict_merge, get_nodes, get_cutoff

from aiida_nanotech_empa.utils import common_utils

Cp2kBaseWorkChain = WorkflowFactory('cp2k.base')

class Cp2kMoleculeOptWorkChain(WorkChain):
    @classmethod
    def define(cls, spec):
        super().define(spec)
        spec.input("code", valid_type=Code)
        spec.input("structure", valid_type=StructureData)

        spec.input(
            "charge",  # +1 means one electron removed
            valid_type=Int,
            default=lambda: Int(0),
            required=False)
        spec.input("multiplicity",
                   valid_type=Int,
                   default=lambda: Int(0),
                   required=False)
        spec.input("magnetization_per_site",
                   valid_type=List,
                   default=lambda: List(list=[]),
                   required=False)
        spec.input("vdw",
                   valid_type=Bool,
                   default=lambda: Bool(False),
                   required=False)
        spec.input("walltime_seconds",
                   valid_type=Int,
                   default=lambda: Int(7200),
                   required=False)
        spec.input("debug",
                   valid_type=Bool,
                   default=lambda: Bool(False),
                   required=False,
                   help="Run with fast parameters for debugging.")

        #workchain outline
        spec.outline(cls.setup, cls.submit_calc, cls.finalize)
        spec.outputs.dynamic = True

        spec.exit_code(
            390,
            "ERROR_TERMINATION",
            message="One or more steps of the work chain failed.",
        )

    def setup(self):
        self.report("Inspecting input and setting up things")

        # --------------------------------------------------

    def submit_calc(self):

        #load input template
        with open(pathlib.Path(__file__).parent /
                  './protocols/molecule_opt_protocol.yml',
                  encoding='utf-8') as handle:
            protocols = yaml.safe_load(handle)
            input_dict = copy.deepcopy(protocols['default'])

        structure = self.inputs.structure
        #cutoff
        self.ctx.cutoff = get_cutoff(structure=structure)

        #get initial magnetization
        magnetization_per_site = copy.deepcopy(
            self.inputs.magnetization_per_site)
        structure_with_tags, kinds_dict = determine_kinds(
            structure, magnetization_per_site)

        #make sure cell is big enough for MT poisson solver
        if self.inputs.debug:
            extra_cell = 5.0
        else:
            extra_cell = 15.0
        self.ctx.atoms = structure_with_tags.get_ase()
        self.ctx.atoms.cell = 2 * (np.ptp(self.ctx.atoms.positions,
                                          axis=0)) + extra_cell
        self.ctx.atoms.center()

        builder = Cp2kBaseWorkChain.get_builder()
        builder.cp2k.code = self.inputs.code
        builder.cp2k.structure = StructureData(ase=self.ctx.atoms)
        builder.cp2k.file = {
            'basis':
            SinglefileData(
                file=os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                  ".", "data", "BASIS_MOLOPT")),
            'pseudo':
            SinglefileData(
                file=os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                  ".", "data", "POTENTIAL")),
        }

        # vdw
        if not self.inputs.vdw.value:
            input_dict['FORCE_EVAL']['DFT']['XC'].pop('VDW_POTENTIAL')

        #UKS
        if self.inputs.multiplicity.value > 0:
            input_dict['FORCE_EVAL']['DFT']['UKS'] = '.TRUE.'
            input_dict['FORCE_EVAL']['DFT'][
                'MULTIPLICITY'] = self.inputs.multiplicity.value

        #cutoff
        input_dict['FORCE_EVAL']['DFT']['MGRID']['CUTOFF'] = self.ctx.cutoff

        if self.inputs.debug:
            input_dict['MOTION']['GEO_OPT']['MAX_FORCE'] = 0.001
            input_dict['FORCE_EVAL']['DFT']['SCF']['EPS_SCF'] = 1e-6
            input_dict['FORCE_EVAL']['DFT']['SCF']['OUTER_SCF'][
                'EPS_SCF'] = 1e-6

        # KINDS section
        self.ctx.kinds_section = get_kinds_section(kinds_dict, protocol='gpw')
        dict_merge(input_dict, self.ctx.kinds_section)

        #computational resources
        nodes, tasks_per_node, threads = get_nodes(
            atoms=self.ctx.atoms,
            calctype='default',
            computer=self.inputs.code.computer,
            max_nodes=48,
            uks=self.inputs.multiplicity.value > 0)

        builder.cp2k.metadata.options.resources = {
            'num_machines': nodes,
            'num_mpiprocs_per_machine': tasks_per_node,
            'num_cores_per_mpiproc': threads
        }

        #walltime
        input_dict['GLOBAL']['WALLTIME'] = max(
            self.inputs.walltime_seconds.value - 600, 600)
        builder.cp2k.metadata.options.max_wallclock_seconds = self.inputs.walltime_seconds.value

        #parser
        builder.cp2k.metadata.options.parser_name = "cp2k_advanced_parser"

        #handlers
        builder.handler_overrides = Dict(
            dict={'resubmit_unconverged_geometry': True})

        #cp2k input dictionary
        builder.cp2k.parameters = Dict(dict=input_dict)

        submitted_node = self.submit(builder)
        return ToContext(opt=submitted_node)

    def finalize(self):
        self.report("Finalizing...")

        if not common_utils.check_if_calc_ok(self, self.ctx.opt):
            return self.exit_codes.ERROR_TERMINATION

        for out in self.ctx.opt.outputs:
            self.out(out, self.ctx.opt.outputs[out])

        # Add extras
        struc = self.ctx.opt.outputs.output_structure
        self.node.set_extra('thumbnail',
                            common_utils.thumbnail(ase_struc=struc.get_ase()))
        self.node.set_extra('formula', struc.get_formula())

        return ExitCode(0)

max_stars_repo_name: nanotech-empa/aiida-nanotech-empa

max_issues_repo_name: nanotech-empa/aiida-nanotech-empa

max_forks_repo_name: nanotech-empa/aiida-nanotech-empa


text:
import numpy as np
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

""" attention pad mask """
def get_attn_pad_mask(seq_q, seq_k, i_pad):
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    pad_attn_mask = seq_k.data.eq(i_pad).unsqueeze(1).expand(batch_size, len_q, len_k)  # <pad>
    return pad_attn_mask

""" attention decoder mask """
def get_attn_decoder_mask(seq):
    subsequent_mask = torch.ones_like(seq).unsqueeze(-1).expand(seq.size(0), seq.size(1), seq.size(1))
    subsequent_mask = subsequent_mask.triu(diagonal=1) # upper triangular part of a matrix(2-D)
    return subsequent_mask

""" scale dot product attention """
class ScaledDotProductAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.dropout = nn.Dropout(config.dropout)
        self.scale = 1 / (self.config.d_head ** 0.5)

    def forward(self, Q, K, V, attn_mask):
        # (bs, n_head, n_q_seq, n_k_seq)
        scores = torch.matmul(Q, K.transpose(-1, -2)).mul_(self.scale)
        scores.masked_fill_(attn_mask, -1e9)
        # (bs, n_head, n_q_seq, n_k_seq)
        attn_prob = nn.Softmax(dim=-1)(scores)
        attn_prob = self.dropout(attn_prob)
        # (bs, n_head, n_q_seq, d_v)
        context = torch.matmul(attn_prob, V)
        # (bs, n_head, n_q_seq, d_v), (bs, n_head, n_q_seq, n_v_seq)
        return context, attn_prob

""" multi head attention """
class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.W_Q = nn.Linear(self.config.d_hidn, self.config.n_head * self.config.d_head)
        self.W_K = nn.Linear(self.config.d_hidn, self.config.n_head * self.config.d_head)
        self.W_V = nn.Linear(self.config.d_hidn, self.config.n_head * self.config.d_head)
        self.scaled_dot_attn = ScaledDotProductAttention(self.config)
        self.linear = nn.Linear(self.config.n_head * self.config.d_head, self.config.d_hidn)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, Q, K, V, attn_mask):
        batch_size = Q.size(0)
        # (bs, n_head, n_q_seq, d_head)
        q_s = self.W_Q(Q).view(batch_size, -1, self.config.n_head, self.config.d_head).transpose(1,2)
        # (bs, n_head, n_k_seq, d_head)
        k_s = self.W_K(K).view(batch_size, -1, self.config.n_head, self.config.d_head).transpose(1,2)
        # (bs, n_head, n_v_seq, d_head)
        v_s = self.W_V(V).view(batch_size, -1, self.config.n_head, self.config.d_head).transpose(1,2)

        # (bs, n_head, n_q_seq, n_k_seq)
        attn_mask = attn_mask.unsqueeze(1).repeat(1, self.config.n_head, 1, 1)

        # (bs, n_head, n_q_seq, d_head), (bs, n_head, n_q_seq, n_k_seq)
        context, attn_prob = self.scaled_dot_attn(q_s, k_s, v_s, attn_mask)
        # (bs, n_head, n_q_seq, h_head * d_head)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.config.n_head * self.config.d_head)
        # (bs, n_head, n_q_seq, e_embd)
        output = self.linear(context)
        output = self.dropout(output)
        # (bs, n_q_seq, d_hidn), (bs, n_head, n_q_seq, n_k_seq)
        return output, attn_prob

""" feed forward """
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.conv1 = nn.Conv1d(in_channels=self.config.d_hidn, out_channels=self.config.d_ff, kernel_size=1)
        self.conv2 = nn.Conv1d(in_channels=self.config.d_ff, out_channels=self.config.d_hidn, kernel_size=1)
        self.active = F.gelu
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, inputs):
        # (bs, d_ff, n_seq)
        output = self.active(self.conv1(inputs.transpose(1, 2)))
        # (bs, n_seq, d_hidn)
        output = self.conv2(output).transpose(1, 2)
        output = self.dropout(output)
        # (bs, n_seq, d_hidn)
        return output

""" encoder layer """
class EncoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.self_attn = MultiHeadAttention(self.config)
        self.layer_norm1 = nn.LayerNorm(self.config.d_hidn, eps=self.config.layer_norm_epsilon)
        self.pos_ffn = PoswiseFeedForwardNet(self.config)
        self.layer_norm2 = nn.LayerNorm(self.config.d_hidn, eps=self.config.layer_norm_epsilon)

    def forward(self, inputs, attn_mask):
        # (bs, n_enc_seq, d_hidn), (bs, n_head, n_enc_seq, n_enc_seq)
        att_outputs, attn_prob = self.self_attn(inputs, inputs, inputs, attn_mask)
        att_outputs = self.layer_norm1(inputs + att_outputs)
        # (bs, n_enc_seq, d_hidn)
        ffn_outputs = self.pos_ffn(att_outputs)
        ffn_outputs = self.layer_norm2(ffn_outputs + att_outputs)
        # (bs, n_enc_seq, d_hidn), (bs, n_head, n_enc_seq, n_enc_seq)
        return ffn_outputs, attn_prob

""" encoder """
class Encoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.enc_emb = nn.Embedding(self.config.n_enc_vocab, self.config.d_embd)
        self.fectorize = nn.Linear(self.config.d_embd, self.config.d_hidn)
        self.pos_emb = nn.Embedding(self.config.n_dec_seq + 1, self.config.d_hidn)
        self.seg_emb = nn.Embedding(self.config.n_seg_type, self.config.d_hidn)

        self.layer = EncoderLayer(self.config)

    def forward(self, inputs, segments):
        positions = torch.arange(inputs.size(1), device=inputs.device, dtype=inputs.dtype).expand(inputs.size(0), inputs.size(1)).contiguous() + 1
        pos_mask = inputs.eq(self.config.i_pad)
        positions.masked_fill_(pos_mask, 0)

        # (bs, n_enc_seq, d_hidn)
        outputs = self.fectorize(self.enc_emb(inputs)) + self.pos_emb(positions)  + self.seg_emb(segments)

        # (bs, n_enc_seq, n_enc_seq)
        attn_mask = get_attn_pad_mask(inputs, inputs, self.config.i_pad)

        attn_probs = []
        for i in range(self.config.n_layer):
            # (bs, n_enc_seq, d_hidn), (bs, n_head, n_enc_seq, n_enc_seq)
            outputs, attn_prob = self.layer(outputs, attn_mask)
            attn_probs.append(attn_prob)

        # (bs, n_enc_seq, d_hidn), [(bs, n_head, n_enc_seq, n_enc_seq)]
        return outputs, attn_probs

""" bert """
class ALBERT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.encoder = Encoder(self.config)

    def forward(self, inputs, segments):
        # (bs, n_seq, d_hidn), [(bs, n_head, n_dec_seq, n_dec_seq)]
        outputs, self_attn_probs = self.encoder(inputs, segments)
        # (bs, n_dec_seq, n_dec_vocab), [(bs, n_head, n_dec_seq, n_dec_seq)]
        return outputs, self_attn_probs

    def save(self, epoch, loss, path):
        torch.save({
            "epoch": epoch,
            "loss": loss,
            "state_dict": self.state_dict()
        }, path)

    def load(self, path):
        save = torch.load(path)
        self.load_state_dict(save["state_dict"])
        return save["epoch"], save["loss"]

""" BERT pretrain """
class ALBERTPretrain(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.albert = ALBERT(self.config)
        # classfier
        self.projection_cls = nn.Linear(self.config.d_hidn, 2, bias=False)
        # lm
        self.linear_lm = nn.Linear(self.config.d_hidn, self.config.d_embd)
        self.projection_lm = nn.Linear(self.config.d_embd, self.config.n_dec_vocab, bias=False)
        self.projection_lm.weight = self.albert.encoder.enc_emb.weight

    def forward(self, inputs, segments):
        # (bs, n_dec_seq, d_hidn), [(bs, n_head, n_dec_seq, n_dec_seq)]
        outputs, attn_probs = self.albert(inputs, segments)
        # (bs, 2)
        logits_cls = self.projection_cls(outputs[:, 0])
        # (bs, n_dec_seq, n_dec_vocab)
        outputs = self.linear_lm(outputs)
        logits_lm = self.projection_lm(outputs)
        # (bs, n_dec_seq - 1, n_enc_vocab), (bs, n_output), [(bs, n_head, n_dec_seq, n_dec_seq)]
        return logits_cls, logits_lm, attn_probs

""" naver movie classfication """
class MovieClassification(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.albert = ALBERT(self.config)
        # classfier
        self.projection_cls = nn.Linear(self.config.d_hidn, self.config.n_output, bias=False)

    def forward(self, inputs, segments):
        # (bs, n_dec_seq, d_hidn), [(bs, n_head, n_dec_seq, n_dec_seq)]
        output, attn_probs = self.albert(inputs, segments)
        # (bs, d_hidn)
        output = output[:, 0]
        # (bs, n_output)
        logits_cls = self.projection_cls(output)
        # (bs, n_output), [(bs, n_head, n_dec_seq, n_dec_seq)]
        return logits_cls, attn_probs

    def save(self, epoch, loss, score, path):
        torch.save({
            "epoch": epoch,
            "loss": loss,
            "score": score,
            "state_dict": self.state_dict()
        }, path)

    def load(self, path):
        save = torch.load(path)
        self.load_state_dict(save["state_dict"])
        return save["epoch"], save["loss"], save["score"]

max_stars_repo_name: kyuhyoung/transformer-evolution

max_issues_repo_name: kyuhyoung/transformer-evolution

max_forks_repo_name: kyuhyoung/transformer-evolution


text:
import datetime
import numpy as np
from nba_matchup import CURRENT_WEEK, START_DATE
import scipy.special as sp

from .util import print_roster

CATS = [
   'FGA', 'FGM', 'FTA', 'FTM', '3PTM', 'PTS', 'REB', 'AST', 'ST', 'BLK', 'TO'
]
NON_PERCENT_CATS = [
   'FGA', 'FTA', '3PTM', 'PTS', 'REB', 'AST', 'ST', 'BLK', 'TO'
]
CATEGORY_NAMES = [
   'FG%', 'FT%', '3PTM', 'PTS', 'REB', 'AST', 'ST', 'BLK', 'TO'
]

IGNORE_POSITIONS = ['BN', 'IL']

def simulate_h2h(roster1, roster2, week=CURRENT_WEEK, num_days=14,
                 num_samples=10000, decay_rate=0.1,
                 include_bench=False,
                 include_injured=False):
    teams = [roster1, roster2]

    diff = (datetime.datetime.today().date() - START_DATE).days // 7
    if week is not None:
        base = START_DATE + datetime.timedelta(days=7 * (week - 1))
    else:
        base = START_DATE + datetime.timedelta(weeks=diff)
    scores, projections = [], []
    week_length = 7 #14 if week == 18 else 7
    ignore = {'BN', 'IL'}
    if include_bench:
        ignore.remove("BN")
    if include_injured:
        ignore.remove("IL")
    for team in teams:
        team_stats, player_games = team.stats(num_days, base_date=base,
                                              week_length=week_length)
        valid_players = set(team_stats[(team_stats["GP"] > 0) &
                                              (include_bench or team_stats["Position"] != "BN")
                                              & (include_injured or team_stats["Position"] !=
                                                 "IL")]['Name'])
        mean, std = compute_average(team_stats, decay_rate=decay_rate)
        num_games = [(player, len(p)) for p, player in zip(player_games,
                                                         team.players)]
        for player, num_game in num_games:
            mean.at[player.name, "Num Games"] = num_game
        for player, num_game in num_games:
            std.at[player.name, "Num Games"] = 0
        score, projection = projected_stats((mean, std), valid_players, num_samples=num_samples)
        valid_index = np.arange(len(mean.index))[mean.index.isin(valid_players)]
        scores.append(score[:, valid_index])
        projections.append(projection)
    cats, points = score_teams(*scores)
    return cats, points, scores, projections

def compute_average(team_stats, decay_rate=0.1):
    team_stats["Weight"] = np.exp(-decay_rate * team_stats["Days Ago"])
    grouped = team_stats.groupby("Name")
    def mean_func(x):
        normed_weights = x['Weight'] / x['Weight'].sum()
        mean = x[CATS].mul(normed_weights, axis=0).sum()
        mean['TFGM'] = mean['FGM'] * len(normed_weights)
        mean['TFGA'] = mean['FGA'] * len(normed_weights)
        mean['TFTM'] = mean['FTM'] * len(normed_weights)
        mean['TFTA'] = mean['FTA'] * len(normed_weights)
        return mean
    def std_func(x):
        data = x[CATS]
        normed_weights = x['Weight'] / x['Weight'].sum()
        mean = data.mul(x["Weight"], axis=0)
        deviation = (data.sub(mean, axis=1) ** 2).mul(x["Weight"], axis=0).sum(axis=0)
        N = deviation.shape[0]
        return np.sqrt(deviation / (N / (N - 1) * x["Weight"].sum()))
    mean = grouped.apply(mean_func)
    std = grouped.apply(std_func)
    return mean, std

def projected_stats(team, valid_players, num_samples=100):
    sample = np.random.normal(loc=np.tile(team[0][NON_PERCENT_CATS].mul(team[0]["Num Games"], axis=0) , [num_samples, 1, 1]),
                              scale=np.tile(team[1][NON_PERCENT_CATS].mul(team[0]["Num Games"], axis=0), [num_samples, 1, 1]))
    projected = team[0].copy()
    projected[NON_PERCENT_CATS] = sample.mean(axis=0)
    ftp = np.random.beta(1. + team[0]['TFTM'], 1. + team[0]['TFTA'] - team[0]['TFTM'],
        [num_samples, len(team[0])])
    fgp = np.random.beta(1. + team[0]['TFGM'], 1. + team[0]['TFGA'] - team[0]['TFGM'],
        [num_samples, len(team[0])])
    fgp[np.isnan(fgp)] = 0.
    ftp[np.isnan(ftp)] = 0.
    fga = np.maximum(sample[..., NON_PERCENT_CATS.index('FGA')].astype(np.int64), 0)
    fta = np.maximum(sample[..., NON_PERCENT_CATS.index('FTA')].astype(np.int64), 0)
    fgm = np.random.binomial(fga, fgp)
    ftm = np.random.binomial(fta, ftp)
    projected['FGM'] = fgm.mean(axis=0)
    projected['FTM'] = ftm.mean(axis=0)
    projected['FG%'] = projected['FGM'] / projected['FGA']
    projected['FT%'] = projected['FTM'] / projected['FTA']
    cols = ['Num Games', 'FGM', 'FGA', 'FG%', 'FTM', 'FTA', 'FT%', '3PTM', 'PTS', 'REB',
        'AST', 'ST', 'BLK', 'TO']
    projected = projected[cols]
    sample = np.concatenate([fgm[..., None], ftm[..., None], sample], axis=-1)
    return sample, projected

def score_teams(team1, team2):
    cats = []
    for i, team in enumerate([team1, team2]):
        team = team.sum(axis=1)
        fg_percent = (team[..., 0] / team[..., 2])
        ft_percent = (team[..., 1] / team[..., 3])
        cats.append(np.concatenate([fg_percent[..., None], ft_percent[..., None], team[..., 4:]], -1))
    cats = np.stack(cats)
    scores = (cats[0][..., :-1] > cats[1][..., :-1]).sum(axis=-1) + (cats[0][..., -1] < cats[1][..., -1]).astype(np.int64)
    return cats, scores

max_stars_repo_name: sharadmv/nba-fantasy-optimize

max_issues_repo_name: sharadmv/nba-fantasy-optimize

max_forks_repo_name: sharadmv/nba-fantasy-optimize

reason: import numpy,import scipy


text:
from pathlib import Path
import numpy as np

class GeoAlignment:

    @staticmethod
    def ExtractAuxData(folder_dir, extension=".[tT][xX][tT]"):
        folder_dir = Path(folder_dir)
        of_a_kind = list(folder_dir.rglob("*" + extension))
        of_jpg = dict((img.with_suffix(""), img) for img in list(folder_dir.rglob("*.[jJ][pP][gG]")))
        of_png = dict((img.with_suffix(""), img) for img in list(folder_dir.rglob("*.[pP][nN][gG]")))
        img_paths = {**of_jpg, **of_png}

        mapping = {}
        file: Path
        for file in of_a_kind:
            base_name = file.with_suffix("")
            if base_name in img_paths:
                mapping[img_paths[base_name]] = file

        return mapping

    @staticmethod
    def CreateGeoHash(point):
        bits = np.unpackbits(point.astype(np.int32).view(np.uint8))
        geo_hash = [None] * len(bits)
        for i in range(0, 32):
            p = i * 3
            geo_hash[p] = bits[i]
            geo_hash[p + 1] = bits[i + 32]
            geo_hash[p + 2] = bits[i + 64]
        return np.packbits(geo_hash[:64])

    @staticmethod
    def EstimatedPoses2Colmap(database_path: Path, estimated_camera_poses):
        raise NotImplementedError("Sorry Quaternion is not know @ this point in time")
    #     conn = COLMAPDatabase.connect(database_path)
    #     images_rows = conn.execute("SELECT * FROM images").fetchall()
    #     cur = conn.cursor()
    #
    #     for row in images_rows:
    #         if row[1] in estimated_camera_poses:
    #             cur.execute("UPDATE licontacts310317 SET liemail=%s, limobile=%s WHERE % s =?" % (liemailval, limobileval, id), (constrain,))
    #
    #
    # S
    #     images_rows = db.execute("SELECT * FROM images").fetchall()
    #     cameras_rows = db.execute("SELECT * FROM cameras").fetchall()
    #     db.close()
    #
    #     images = {}
    #     for row in images_rows:
    #         old = rehash[row[1]]
    #         images[row[0]] = BaseImage(row[0], old.qvec, old.tvec, row[2], old.name, [], [])
    #     model.images = images
    #
    #     cameras = {}
    #     for row in cameras_rows:
    #         camera_id, cam_model, width, height, params, prior = row
    #         params = blob_to_array(params, np.float64)
    #         cameras[row[0]] = Camera(camera_id, CAMERA_MODEL_IDS[cam_model].model_name, width, height, params)
    #     model.cameras = cameras

    # return model

# if __name__ == "__main__":
#     proj = config.SparseModel(config.project_path, model_path=config.project_path + "/sparse/0")
#     GeoAlignment.Dlr2Colmap(proj.images_path, proj.base_path, ECEF=False)


text:
"""
Gym env wrappers that make the environment suitable for the RL algorithms.

"""
import datetime
import json
import os
from collections import deque
from os.path import join

import cv2
import gym
import numpy as np
# noinspection PyProtectedMember
from gym import spaces, RewardWrapper, ObservationWrapper

from seed_rl.algorithms.utils.algo_utils import num_env_steps
from seed_rl.utils.utils import numpy_all_the_way, ensure_dir_exists, log

def reset_with_info(env):
    """Sometimes we want to get info with the very first frame."""
    obs = env.reset()
    info = {}
    if hasattr(env.unwrapped, 'get_info_all'):
        info = env.unwrapped.get_info_all()  # info for the new episode
    return obs, info

def unwrap_env(wrapped_env):
    return wrapped_env.unwrapped

def is_goal_based_env(env):
    dict_obs = isinstance(env.observation_space, spaces.Dict)
    if not dict_obs:
        return False

    for key in ['obs', 'goal']:
        if key not in env.observation_space.spaces:
            return False

    return True

def main_observation_space(env):
    if hasattr(env.observation_space, 'spaces'):
        return env.observation_space.spaces['obs']
    else:
        return env.observation_space

def has_image_observations(observation_space):
    """It's a heuristic."""
    return len(observation_space.shape) >= 2

class StackFramesWrapper(gym.core.Wrapper):
    """
    Gym env wrapper to stack multiple frames.
    Useful for training feed-forward agents on dynamic games.
    """

    def __init__(self, env, stack_past_frames, channel_config='HWC'):
        super(StackFramesWrapper, self).__init__(env)
        if len(env.observation_space.shape) not in [1, 2]:
            raise Exception('Stack frames works with vector observations and 2D single channel images')
        self._stack_past = stack_past_frames
        self._frames = None

        self._image_obs = has_image_observations(env.observation_space)

        self.channel_config = channel_config
        if self._image_obs:
            if self.channel_config == 'CHW':
                new_obs_space_shape = (stack_past_frames,) + env.observation_space.shape
            elif self.channel_config == 'HWC':
                new_obs_space_shape = env.observation_space.shape + (stack_past_frames,)
            else:
                raise Exception(f'Unknown channel config {self.channel_config}')
        else:
            new_obs_space_shape = list(env.observation_space.shape)
            new_obs_space_shape[0] *= stack_past_frames

        self.observation_space = spaces.Box(
            env.observation_space.low.flat[0],
            env.observation_space.high.flat[0],
            shape=new_obs_space_shape,
            dtype=env.observation_space.dtype,
        )

    def _render_stacked_frames(self):
        if self._image_obs:
            # stack past frames along first dimension
            img = numpy_all_the_way(self._frames)

            if self.channel_config == 'CHW':
                return img
            elif self.channel_config == 'HWC':
                return np.transpose(img, axes=[1, 2, 0])
            else:
                raise Exception(f'Unknown channel config {self.channel_config}')
        else:
            return np.array(self._frames).flatten()

    def reset(self):
        observation = self.env.reset()
        self._frames = deque([observation] * self._stack_past)
        return self._render_stacked_frames()

    def step(self, action):
        new_observation, reward, done, info = self.env.step(action)
        self._frames.popleft()
        self._frames.append(new_observation)
        return self._render_stacked_frames(), reward, done, info

class SkipFramesWrapper(gym.core.Wrapper):
    """Wrapper for action repeat over N frames to speed up training."""

    def __init__(self, env, skip_frames=4):
        super(SkipFramesWrapper, self).__init__(env)
        self._skip_frames = skip_frames

    def reset(self):
        return self.env.reset()

    def step(self, action):
        done = False
        info = None
        new_observation = None

        total_reward, num_frames = 0, 0
        for i in range(self._skip_frames):
            new_observation, reward, done, info = self.env.step(action)
            num_frames += 1
            total_reward += reward
            if done:
                break

        info['num_frames'] = num_frames
        return new_observation, total_reward, done, info

class SkipAndStackFramesWrapper(StackFramesWrapper):
    """Wrapper for action repeat + stack multiple frames to capture dynamics."""

    def __init__(self, env, skip_frames=4, stack_frames=4, channel_config='HWC'):
        super().__init__(env, stack_past_frames=stack_frames, channel_config=channel_config)
        self._skip_frames = skip_frames

    def step(self, action):
        done = False
        info = {}
        total_reward, num_frames = 0, 0
        for i in range(self._skip_frames):
            new_observation, reward, done, info = self.env.step(action)
            num_frames += 1
            total_reward += reward
            self._frames.popleft()
            self._frames.append(new_observation)
            if done:
                break

        info['num_frames'] = num_frames
        return self._render_stacked_frames(), total_reward, done, info

class NormalizeWrapper(gym.core.Wrapper):
    """
    For environments with vector lowdim input.

    """

    def __init__(self, env):
        super(NormalizeWrapper, self).__init__(env)
        if len(env.observation_space.shape) != 1:
            raise Exception('NormalizeWrapper only works with lowdimensional envs')

        self.wrapped_env = env
        self._normalize_to = 1.0

        self._mean = (env.observation_space.high + env.observation_space.low) * 0.5
        self._max = env.observation_space.high

        self.observation_space = spaces.Box(
            -self._normalize_to, self._normalize_to, shape=env.observation_space.shape, dtype=np.float32,
        )

    def _normalize(self, obs):
        obs -= self._mean
        obs *= self._normalize_to / (self._max - self._mean)
        return obs

    def reset(self):
        observation = self.env.reset()
        return self._normalize(observation)

    def step(self, action):
        observation, reward, done, info = self.env.step(action)
        return self._normalize(observation), reward, done, info

    @property
    def range(self):
        return [-self._normalize_to, self._normalize_to]

class ResizeWrapper(gym.core.Wrapper):
    """Resize observation frames to specified (w,h) and convert to grayscale."""

    def __init__(self, env, w, h, grayscale=True, add_channel_dim=False, area_interpolation=False):
        super(ResizeWrapper, self).__init__(env)

        self.w = w
        self.h = h
        self.grayscale = grayscale
        self.add_channel_dim = add_channel_dim
        self.interpolation = cv2.INTER_AREA if area_interpolation else cv2.INTER_NEAREST

        if isinstance(env.observation_space, spaces.Dict):
            # TODO: does this even work?
            new_spaces = {}
            for key, space in env.observation_space.spaces.items():
                new_spaces[key] = self._calc_new_obs_space(space)
            self.observation_space = spaces.Dict(new_spaces)
        else:
            self.observation_space = self._calc_new_obs_space(env.observation_space)

    def _calc_new_obs_space(self, old_space):
        low, high = old_space.low.flat[0], old_space.high.flat[0]

        if self.grayscale:
            new_shape = [self.h, self.w, 1] if self.add_channel_dim else [self.h, self.w]
        else:
            if len(old_space.shape) > 2:
                channels = old_space.shape[-1]
                new_shape = [self.h, self.w, channels]
            else:
                new_shape = [self.h, self.w, 1] if self.add_channel_dim else [self.h, self.w]

        return spaces.Box(low, high, shape=new_shape, dtype=old_space.dtype)

    def _convert_obs(self, obs):
        if obs is None:
            return obs

        obs = cv2.resize(obs, (self.w, self.h), interpolation=self.interpolation)
        if self.grayscale:
            obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)

        if self.add_channel_dim:
            return obs[:, :, None]  # add new dimension (expected by tensorflow)
        else:
            return obs

    def _observation(self, obs):
        if isinstance(obs, dict):
            new_obs = {}
            for key, value in obs.items():
                new_obs[key] = self._convert_obs(value)
            return new_obs
        else:
            return self._convert_obs(obs)

    def reset(self):
        return self._observation(self.env.reset())

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        return self._observation(obs), reward, done, info

class VerticalCropWrapper(ObservationWrapper):
    def __init__(self, env, crop_h):
        # super(VerticalCropWrapper, self).__init__(env)
        super().__init__(env)

        self.crop_h = crop_h
        self.observation_space = self._calc_new_obs_space(env.observation_space)

    def _calc_new_obs_space(self, old_space):
        low, high = old_space.low.flat[0], old_space.high.flat[0]
        h, w, channels = old_space.shape
        new_shape = [self.crop_h, w, channels]
        return spaces.Box(low, high, shape=new_shape, dtype=old_space.dtype)

    # noinspection PyProtectedMember
    def observation(self, observation):
        h = observation.shape[0]
        crop_top = (h - self.crop_h) // 2
        crop_bottom = h - self.crop_h - crop_top
        cropped_obs = observation[crop_top:h - crop_bottom, :, :]
        return cropped_obs

class RewardScalingWrapper(RewardWrapper):
    def __init__(self, env, scaling_factor):
        super(RewardScalingWrapper, self).__init__(env)
        self._scaling = scaling_factor
        self.reward_range = (r * scaling_factor for r in self.reward_range)

    def reward(self, reward):
        return reward * self._scaling

class TimeLimitWrapper(gym.core.Wrapper):
    terminated_by_timer = 'terminated_by_timer'

    def __init__(self, env, limit, random_variation_steps=0):
        super(TimeLimitWrapper, self).__init__(env)
        self._limit = limit
        self._variation_steps = random_variation_steps
        self._num_steps = 0
        self._terminate_in = self._random_limit()

    def _random_limit(self):
        return np.random.randint(-self._variation_steps, self._variation_steps + 1) + self._limit

    def reset(self):
        self._num_steps = 0
        self._terminate_in = self._random_limit()
        return self.env.reset()

    def step(self, action):
        observation, reward, done, info = self.env.step(action)
        if observation is None:
            return observation, reward, done, info

        self._num_steps += num_env_steps([info])
        if done:
            pass
        else:
            if self._num_steps >= self._terminate_in:
                done = True
                info[self.terminated_by_timer] = True

        return observation, reward, done, info

class RemainingTimeWrapper(ObservationWrapper):
    """Designed to be used together with TimeLimitWrapper."""

    def __init__(self, env):
        super(RemainingTimeWrapper, self).__init__(env)

        # adding an additional input dimension to indicate time left before the end of episode
        self.observation_space = spaces.Dict({
            'timer': spaces.Box(0.0, 1.0, shape=[1], dtype=np.float32),
            'obs': env.observation_space,
        })

        wrapped_env = env
        while not isinstance(wrapped_env, TimeLimitWrapper):
            wrapped_env = wrapped_env.env
            if not isinstance(wrapped_env, gym.core.Wrapper):
                raise Exception('RemainingTimeWrapper is supposed to wrap TimeLimitWrapper')
        self.time_limit_wrapper = wrapped_env

    # noinspection PyProtectedMember
    def observation(self, observation):
        num_steps = self.time_limit_wrapper._num_steps
        terminate_in = self.time_limit_wrapper._terminate_in

        dict_obs = {
            'timer': num_steps / terminate_in,
            'obs': observation,
        }
        return dict_obs

class PixelFormatChwWrapper(ObservationWrapper):
    """TODO? This can be optimized for VizDoom, can we query CHW directly from VizDoom?"""

    def __init__(self, env):
        super().__init__(env)

        if isinstance(env.observation_space, gym.spaces.Dict):
            img_obs_space = env.observation_space['obs']
            self.dict_obs_space = True
        else:
            img_obs_space = env.observation_space
            self.dict_obs_space = False

        if not has_image_observations(img_obs_space):
            raise Exception('Pixel format wrapper only works with image-based envs')

        obs_shape = img_obs_space.shape
        max_num_img_channels = 4

        if len(obs_shape) <= 2:
            raise Exception('Env obs do not have channel dimension?')

        if obs_shape[0] <= max_num_img_channels:
            raise Exception('Env obs already in CHW format?')

        h, w, c = obs_shape
        low, high = img_obs_space.low.flat[0], img_obs_space.high.flat[0]
        new_shape = [c, h, w]

        if self.dict_obs_space:
            dtype = env.observation_space.spaces['obs'].dtype if env.observation_space.spaces['obs'].dtype is not None else np.float32
        else:
            dtype = env.observation_space.dtype if env.observation_space.dtype is not None else np.float32

        new_img_obs_space = spaces.Box(low, high, shape=new_shape, dtype=dtype)

        if self.dict_obs_space:
            self.observation_space = env.observation_space
            self.observation_space.spaces['obs'] = new_img_obs_space
        else:
            self.observation_space = new_img_obs_space

        self.action_space = env.action_space

    @staticmethod
    def _transpose(obs):
        return np.transpose(obs, (2, 0, 1))  # HWC to CHW for PyTorch

    def observation(self, observation):
        if observation is None:
            return observation

        if self.dict_obs_space:
            observation['obs'] = self._transpose(observation['obs'])
        else:
            observation = self._transpose(observation)
        return observation

class ClipRewardWrapper(gym.RewardWrapper):
    def __init__(self, env):
        gym.RewardWrapper.__init__(self, env)

    def reward(self, reward):
        reward = min(5.0, reward)
        reward = max(-0.1, reward)
        return reward

class RecordingWrapper(gym.core.Wrapper):
    def __init__(self, env, record_to, player_id):
        super().__init__(env)

        self._record_to = record_to
        self._episode_recording_dir = None
        self._record_id = 0
        self._frame_id = 0
        self._player_id = player_id
        self._recorded_episode_reward = 0
        self._recorded_episode_shaping_reward = 0

        self._recorded_actions = []

        # Experimental! Recording Doom replay. Does not work in all scenarios, e.g. when there are in-game bots.
        self.unwrapped.record_to = record_to

    def reset(self):
        if self._episode_recording_dir is not None and self._record_id > 0:
            # save actions to text file
            with open(join(self._episode_recording_dir, 'actions.json'), 'w') as actions_file:
                json.dump(self._recorded_actions, actions_file)

            # rename previous episode dir
            reward = self._recorded_episode_reward + self._recorded_episode_shaping_reward
            new_dir_name = self._episode_recording_dir + f'_r{reward:.2f}'
            os.rename(self._episode_recording_dir, new_dir_name)
            log.info(
                'Finished recording %s (rew %.3f, shaping %.3f)',
                new_dir_name, reward, self._recorded_episode_shaping_reward,
            )

        dir_name = f'ep_{self._record_id:03d}_p{self._player_id}'
        self._episode_recording_dir = join(self._record_to, dir_name)
        ensure_dir_exists(self._episode_recording_dir)

        self._record_id += 1
        self._frame_id = 0
        self._recorded_episode_reward = 0
        self._recorded_episode_shaping_reward = 0

        self._recorded_actions = []

        return self.env.reset()

    def _record(self, img):
        frame_name = f'{self._frame_id:05d}.png'
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(join(self._episode_recording_dir, frame_name), img)
        self._frame_id += 1

    def step(self, action):
        observation, reward, done, info = self.env.step(action)

        if isinstance(action, np.ndarray):
            self._recorded_actions.append(action.tolist())
        elif isinstance(action, np.int64):
            self._recorded_actions.append(int(action))
        else:
            self._recorded_actions.append(action)

        self._record(observation)
        self._recorded_episode_reward += reward
        if hasattr(self.env.unwrapped, '_total_shaping_reward'):
            # noinspection PyProtectedMember
            self._recorded_episode_shaping_reward = self.env.unwrapped._total_shaping_reward

        return observation, reward, done, info