text:
# ##########################################################
#    FILENAME:   Hydraulic_Property_Calculation.py
#    VERSION:    2.1b
#    SINCE:      2016-08-09
#    AUTHOR:     Xing Zheng - zhengxing@utexas.edu
#    Description:This program is designed for evaluating
#    NHD Flowline Hydraulic Properties from HAND Raster.
#    The input should be (1) Catchment shapefile;
#                        (2) Flowline shapefile;
#                        (3) HAND raster generated
#                            from NHD-HAND.
# ##########################################################

import sys
from osgeo import gdal, ogr
import os
import numpy as np
import numpy.ma as ma
from math import sqrt
import shutil

def Shapefile_Attribute_Reader(catchmentshp, flowlineshp):
    """Find the shared COMID/FEATURE between flowline feature class
    and catchment feature class and read some flowline attributes
    needed to calculate channel hydraulic properties
    """
    # Read catchment shapefile
    driver = ogr.GetDriverByName("ESRI Shapefile")
    dataSource = driver.Open(catchmentshp, 0)
    layer = dataSource.GetLayer()
    # Get catchment FEATUREID list
    catchment_COMID = []
    for catchment in layer:
        catchment_COMID.append(catchment.GetField("FEATUREID"))
    # Get flowline COMID list, reach length, and slope
    flowline_COMID = []
    RiverLength_dic = {}
    Slope_dic = {}
    dataSource = driver.Open(flowlineshp, 0)
    layer = dataSource.GetLayer()
    for flowline in layer:
        flowline_COMID.append(flowline.GetField("COMID"))
        RiverLength_dic[str(flowline.GetField("COMID"))] = flowline.GetField("LENGTHKM")
        Slope_dic[str(flowline.GetField("COMID"))] = flowline.GetField("SLOPE")
    # Find the intersection between catchment FEATUREID set
    # and flowline COMID set
    COMIDlist = list(set(catchment_COMID).intersection(flowline_COMID))

    return COMIDlist, RiverLength_dic, Slope_dic

def HANDClipper(catchmentshp, flowlineshp, handtif,
                COMIDlist, RiverLength_dic, Slope_dic,
                Hmax, dh, roughness):
    """Create HAND Raster for every catchment in the study area
    """
    # Set output path
    outputFolder = os.path.join(os.getcwd(), "HydraulicProperty")
    catchmentFolder = os.path.join(os.getcwd(), "Catchment")
    # Clean up existing folders and files
    if os.path.exists(catchmentFolder):
        shutil.rmtree(catchmentFolder)
    if os.path.exists(outputFolder):
        shutil.rmtree(outputFolder)
    # Create output folder
    if not os.path.exists(catchmentFolder):
        os.mkdir(catchmentFolder)
    if not os.path.exists(outputFolder):
        os.mkdir(outputFolder)
    volumeFolder = os.path.join(outputFolder, "Volume")
    if not os.path.exists(volumeFolder):
        os.mkdir(volumeFolder)
    topwidthFolder = os.path.join(outputFolder, "TopWidth")
    if not os.path.exists(topwidthFolder):
        os.mkdir(topwidthFolder)
    wetareaFolder = os.path.join(outputFolder, "WetArea")
    if not os.path.exists(wetareaFolder):
        os.mkdir(wetareaFolder)
    wetperimeterFolder = os.path.join(outputFolder, "WettedPerimeter")
    if not os.path.exists(wetperimeterFolder):
        os.mkdir(wetperimeterFolder)
    surfaceareaFolder = os.path.join(outputFolder, "SurfaceArea")
    if not os.path.exists(surfaceareaFolder):
        os.mkdir(surfaceareaFolder)
    bedareaFolder = os.path.join(outputFolder, "BedArea")
    if not os.path.exists(bedareaFolder):
        os.mkdir(bedareaFolder)
    hydraulicradiusFolder = os.path.join(outputFolder, "HydraulicRadius")
    if not os.path.exists(hydraulicradiusFolder):
        os.mkdir(hydraulicradiusFolder)
    ratingcurveFolder = os.path.join(outputFolder, "RatingCurve")
    if not os.path.exists(ratingcurveFolder):
        os.mkdir(ratingcurveFolder)
    summaryFolder = os.path.join(outputFolder, "Summary")
    if not os.path.exists(summaryFolder):
        os.mkdir(summaryFolder)
    # Read catchment shapefile
    inDriver = ogr.GetDriverByName("ESRI Shapefile")
    inDataSource = inDriver.Open(catchmentshp, 0)
    inLayer = inDataSource.GetLayer()
    spatialRef = inLayer.GetSpatialRef()
    for i in range(0, inLayer.GetFeatureCount()):
        # Get an input catchment feature
        inFeature = inLayer.GetFeature(i)
        COMID = inFeature.GetField('FeatureID')
        if COMID in COMIDlist:
            # Set individual catchment shapefile path
            outShapefile = os.path.join(catchmentFolder, str(COMID)+".shp")
            outDriver = ogr.GetDriverByName("ESRI Shapefile")
            # Remove output shapefile if it already exists
            if os.path.exists(outShapefile):
                outDriver.DeleteDataSource(outShapefile)
            # Create the output shapefile
            outDataSource = outDriver.CreateDataSource(outShapefile)
            outLayer = outDataSource.CreateLayer(str(COMID), spatialRef,
                                                 geom_type=ogr.wkbPolygon)
            # Add input layer fields to the output Layer
            inLayerDefn = inLayer.GetLayerDefn()
            for k in range(0, inLayerDefn.GetFieldCount()):
                fieldDefn = inLayerDefn.GetFieldDefn(k)
                outLayer.CreateField(fieldDefn)
            # Get the output layer's feature definition
            outLayerDefn = outLayer.GetLayerDefn()
            outFeature = ogr.Feature(outLayerDefn)
            # Add field values from input layer
            for m in range(0, outLayerDefn.GetFieldCount()):
                outFeature.SetField(outLayerDefn.GetFieldDefn(m).GetNameRef(),
                                    inFeature.GetField(m))
            # Get polygon geometry
            geom = inFeature.GetGeometryRef()
            outFeature.SetGeometry(geom)
            # Add new feature to output Layer
            outLayer.CreateFeature(outFeature)
            # Close output file
            outDataSource.Destroy()
            # Read single catchment shapefile
            infile = outShapefile
            inSource = inDriver.Open(infile, 0)
            inlayer = inSource.GetLayer()
            extent = inlayer.GetExtent()
            dtsdir = handtif
            dts_out_file = os.path.join(catchmentFolder,
                                        str(COMID) + "hand.tif")
            if os.path.exists(dts_out_file):
                os.remove(dts_out_file)
            # Clip HAND raster with single catchment polygon boundary
            command_dd = "gdalwarp -te " + str(extent[0]) + " " + \
                         str(extent[2]) + " " + str(extent[1]) + " " + \
                         str(extent[3]) + " -dstnodata -32768 -cutline " + \
                         outShapefile + " -cl " + str(COMID) + " " + dtsdir + \
                         " " + dts_out_file
            os.system(command_dd)
            if os.path.exists(dts_out_file):
                # Hydraulic property calculation
                volumefolder = os.path.join(outputFolder, "Volume")
                volume_write = open(volumefolder+"/"+str(COMID)+".csv",
                                    'a')
                volume_write.write("StageHeight(m),Volume(m^3)\n")
                volume_write.close()
                # Calculate flood volume and water surface area
                Return = Volume_SA_Calculation(COMID, dts_out_file,
                                               dh, Hmax)
                Volume = np.asarray(Return[0])
                SAlist = np.asarray(Return[1])
                for H in np.arange(0, Hmax, dh):
                    volume_write = open(volumefolder+"/"+str(COMID)+".csv",
                                        'a')
                    volume_write.write(str(H)+","+str(Volume[int(H/dh)])+"\n")
                    volume_write.close()
                Depth, Volume = np.loadtxt(volumeFolder + "/" +
                                           str(COMID) + ".csv",
                                           delimiter=',', skiprows=1,
                                           usecols=(0, 1, ), unpack=True)
                # Calculate channel top width, wet area,
                # wetted perimeter, and bed area
                if np.any(Volume):
                    RiverLength = RiverLength_dic[str(COMID)]
                    Return_Result = TW_WA_WP_BA_Calculation(Depth, RiverLength,
                                                            Volume, SAlist, dh)
                    TWlist = np.asarray(Return_Result[0])
                    WAlist = np.asarray(Return_Result[1])
                    WPlist = np.asarray(Return_Result[2])
                    BAlist = np.asarray(Return_Result[3])
                    TWlocation = os.path.join(topwidthFolder,
                                              str(COMID) + '.csv')
                    np.savetxt(TWlocation, np.column_stack((Depth, TWlist)),
                               delimiter=',',
                               header="Depth(m),TopWidth(m)",
                               comments='')
                    WAlocation = os.path.join(wetareaFolder,
                                              str(COMID)+".csv")
                    np.savetxt(WAlocation, np.column_stack((Depth, WAlist)),
                               delimiter=',',
                               header="Depth(m),WetArea(m^2)",
                               comments='')
                    WPlocation = os.path.join(wetperimeterFolder,
                                              str(COMID)+".csv")
                    np.savetxt(WPlocation, np.column_stack((Depth, WPlist)),
                               delimiter=',',
                               header="Depth(m),WettedPerimeter(m)",
                               comments='')
                    SAlocation = os.path.join(surfaceareaFolder,
                                              str(COMID)+".csv")
                    np.savetxt(SAlocation, np.column_stack((Depth, SAlist)),
                               delimiter=',',
                               header="Depth(m),SurfaceArea(m^2)",
                               comments='')
                    BAlocation = os.path.join(bedareaFolder,
                                              str(COMIDlist[i])+".csv")
                    np.savetxt(BAlocation, np.column_stack((Depth, BAlist)),
                               delimiter=',', header="Depth(m),BedArea(m^2)",
                               comments='')
                    # Calculate hydraulic radius and discharge
                    Slope = Slope_dic[str(COMID)]
                    Return_Result = HR_Q_Calculation(WAlist, WPlist,
                                                     Slope, roughness)
                    HRlist = Return_Result[0]
                    Qlist = Return_Result[1]
                    HRlocation = os.path.join(hydraulicradiusFolder,
                                              str(COMID)+".csv")
                    np.savetxt(HRlocation, np.column_stack((Depth, HRlist)),
                               delimiter=',',
                               header="Depth(m),HydraulicRadius(m)",
                               comments='')
                    HQlocation = os.path.join(ratingcurveFolder,
                                              str(COMID)+".csv")
                    np.savetxt(HQlocation, np.column_stack((Depth, Qlist)),
                               delimiter=',',
                               header="Depth(m),Discharge(m^3/s)",
                               comments='')
                    Sumlocation = os.path.join(summaryFolder,
                                               str(COMID)+".csv")
                    np.savetxt(Sumlocation, np.column_stack((Depth, TWlist,
                                                             WAlist, WPlist,
                                                             SAlist, BAlist,
                                                             HRlist, Qlist)),
                               delimiter=',', header="Depth(m),TopWidth(m),"
                               "WetArea(m^2),WettedPerimeter(m),"
                               "SurfaceArea(m^2),BedArea(m^2),"
                               "HydraulicRadius(m),Discharge(m^3/s)",
                               comments='')

def Volume_SA_Calculation(COMID, HANDRaster, dh, Hmax):
    """Calculate flood volume and water surface area
    """
    # Read HAND raster
    dts_ds = gdal.Open(HANDRaster)
    band_dts = dts_ds.GetRasterBand(1)
    nodata_dts = band_dts.GetNoDataValue()
    array_dts = band_dts.ReadAsArray()
    arraydts = ma.masked_where(array_dts == nodata_dts, array_dts)
    Volumelist = []
    SAlist = []
    for H in np.arange(0, Hmax, dh):
        # Subtract stage height from hand raster
        dts_value = arraydts-H
        # Find dts<0
        dts_less_height = dts_value <= 0
        # Count number of cell has negative value
        count_cell = dts_less_height.sum()
        cell_height = dts_value[dts_less_height]*(-1)
        # Calculate flood volume for stage height H
        volume_in = 10*10*cell_height
        volume = volume_in.sum()
        # Calculate water surface area for stage height H
        SurfaceArea = count_cell*10*10
        if type(volume) is not np.float32:
            Volumelist.append(0.0)
            SAlist.append(0.0)
        else:
            Volumelist.append(volume)
            SAlist.append(SurfaceArea)
    ReturnResult = [Volumelist, SAlist]
    return ReturnResult

def TW_WA_WP_BA_Calculation(Depth, RiverLength, Volume, SAlist, dh):
    """Calculate channel top width, wet area, wetted perimeter, and
    bed area
    """
    Volume = Volume - Volume[0]
    DVolume = np.diff(Volume)
    DDepth = np.diff(Depth)
    TotalArea = Volume/RiverLength/1000
    TWlist = SAlist/RiverLength/1000
    WAlist = list(TotalArea)
    TWlist = list(TWlist)
    WPlist = []
    BAlist = []
    WetPerimeter = 0
    BedArea = 0
    for i in range(len(TWlist)):
        if i == 0:
            WetPerimeter = TWlist[i]
            WPlist.append(WetPerimeter)
        else:
            WetPerimeter += 2*sqrt(DDepth[i-1]*DDepth[i-1] +
                                   ((TWlist[i]-TWlist[i-1])/2)**2)
            WPlist.append(WetPerimeter)
        BedArea = WetPerimeter*RiverLength
        BAlist.append(BedArea)
    ReturnResult = [TWlist, WAlist, WPlist, BAlist]

    return ReturnResult

def HR_Q_Calculation(WAlist, WPlist, Slope, roughness):
    """ Calculate hydraulic radius and discharge
    """
    HRlist = []
    Qlist = []
    HRlist.append(0)
    for i in range(1, len(WAlist)):
        HydraulicRadius = WAlist[i]/WPlist[i]
        HRlist.append(HydraulicRadius)
    Qlist.append(0)
    for i in range(1, len(WAlist)):
        if Slope >= 0:
            Discharge = WAlist[i]*(HRlist[i]**(2.0/3))*sqrt(Slope)/roughness
            Qlist.append(Discharge)
        else:
            Qlist.append(0)
    ReturnResult = [HRlist, Qlist]

    return ReturnResult

def main():
    catchmentshp = str(sys.argv[2])
    flowlineshp = str(sys.argv[4])
    handtif = str(sys.argv[6])
    Hmax = 25
    dh = 0.3048
    roughness = 0.05
    COMIDlist, RiverLength_dic, Slope_dic = Shapefile_Attribute_Reader(catchmentshp, flowlineshp)
    HANDClipper(catchmentshp, flowlineshp, handtif,
                COMIDlist, RiverLength_dic, Slope_dic,
                Hmax, dh, roughness)

if __name__ == "__main__":
    main()

max_stars_repo_name: cybergis/nfie-inunmap

max_issues_repo_name: cybergis/nfie-inunmap

max_forks_repo_name: cybergis/nfie-inunmap


text:
"""
Bayesian Aerosol Retrieval algorithm
version 1.01 (12 August 2020)

For the most recent version of the codes, please see https://github.com/TUT-ISI/BARalgorithm

Reference:
Lipponen, A., Mielonen, T., Pitkänen, M. R. A., Levy, R. C., Sawyer, V. R., Romakkaniemi, S., Kolehmainen, V., and Arola, A.:
Bayesian Aerosol Retrieval Algorithm for MODIS AOD retrieval over land, Atmos. Meas. Tech., 11, 1529–1547, 2018.
https://doi.org/10.5194/amt-11-1529-2018.

Contact information:
Antti Lipponen
Finnish Meteorological Institute
antti.lipponen@fmi.fi

----

Copyright 2018-2020 Antti Lipponen / Finnish Meteorological Institute

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation
files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy,
modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""

import os
import json
from scipy.interpolate import RegularGridInterpolator
from netCDF4 import Dataset
import numpy as np
from scipy.sparse import diags, csr_matrix
import sys
from datetime import datetime
# BAR functions
from BARfunctions import loadGranule, prepareModels, loadSurfaceReflectancePrior, combinePolyLookupList, BARretrieve

def printUse():
    print('Use: "python BAR.py INPUT [OUTPUT] [-quantifyuncertainty] [-nospatialcorrelation] [-noapproximationerror]"\nExample: "python BAR.py MOD04_L2.A2015101.1155.006.2015104000600.hdf"')
    print('\nCommand line options:\n-quantifyuncertainty: Quantify the uncertainties related to retrieval and save the uncertainty estimates into the outputfile')
    print('-nospatialcorrelation: Do not use spatial correlation model')
    print('-noapproximationerror: Do not use approximation error model\n')

# Read granule to be retrieved
if len(sys.argv) < 2:
    print('\nMODIS M(O|Y)D04_L2 granule filename needed.')
    printUse()
    sys.exit(0)

dataFile = sys.argv[1]
inputFilePath, inputFileName = os.path.split(dataFile)
print('\nBayesian Aerosol Retrieval (BAR) algorithm version 1.01 (12 August 2020).\n\n  ---------------------\n  Reference:\n  Lipponen, A., Mielonen, T., Pitkänen, M. R. A., Levy, R. C., Sawyer, V. R., Romakkaniemi, S., Kolehmainen, V., and Arola, A.:\n  Bayesian Aerosol Retrieval Algorithm for MODIS AOD retrieval over land, Atmos. Meas. Tech., 11, 1529–1547, 2018. https://doi.org/10.5194/amt-11-1529-2018\n  ---------------------\n  Contact information:\n  Antti Lipponen, Finnish Meteorological Institute\n  antti.lipponen@fmi.fi\n  ---------------------\n')
print('  Retrieving MODIS granule file: {}'.format(inputFileName))

if len(sys.argv) >= 3 and sys.argv[2][0] != '-':
    outputFilePath, outputFileName = os.path.split(sys.argv[2])
else:
    outputFilePath, outputFileName = inputFilePath, 'BAR_' + '.'.join(os.path.split(dataFile)[1].split('.')[:-1]) + '.nc'

if not os.path.isfile(os.path.join(inputFilePath, inputFileName)):
    print('\nMODIS granule file "{}" not found.\n'.format(os.path.join(inputFilePath, inputFileName)))
    sys.exit(1)

# Define areas, prior models & settings
areas = {
    # East North America
    'ENA': {'min_lat': 25.0, 'max_lat': 90.0, 'min_lon': -100.0, 'max_lon': -30.0},
    # West North America
    'WNA': {'min_lat': 25.0, 'max_lat': 90.0, 'min_lon': -180.0, 'max_lon': -100.0},
    # Central and South America
    'CSA': {'min_lat': -90.0, 'max_lat': 25.0, 'min_lon': -180.0, 'max_lon': -30.0},
    # Europe
    'EUR': {'min_lat': 35.0, 'max_lat': 90.0, 'min_lon': -30.0, 'max_lon': 60.0},
    # North Africa and Middle East
    'NAME': {'min_lat': 0.0, 'max_lat': 35.0, 'min_lon': -30.0, 'max_lon': 60.0},
    # South Africa
    'SA': {'min_lat': -90.0, 'max_lat': 0.0, 'min_lon': -30.0, 'max_lon': 60.0},
    # Northeast Asia
    'NEA': {'min_lat': 35.0, 'max_lat': 90.0, 'min_lon': 60.0, 'max_lon': 180.0},
    # Southeast Asia
    'SEA': {'min_lat': 0.0, 'max_lat': 35.0, 'min_lon': 60.0, 'max_lon': 180.0},
    # Oceania
    'OCE': {'min_lat': -90.0, 'max_lat': 0.0, 'min_lon': 60.0, 'max_lon': 180.0}
}

settings = {
    'max_pixels_per_part': 3000,  # max number of pixels to be retrieved at once
    'BFGSftol': 1.0e-6,  # iteration parameter
    'BFGSgtol': 1.0e-6,  # iteration parameter
    'BFGSmaxiter': 5000,  # iteration parameter
    'N_processes': 4,  # number of processes for retrieval (if computer has 4 CPU cores, 4 should result in optimal performance)
    'surfacePriorYear': 2014,  # year of surface reflectance prior model
    'useApproximationErrorModel': True,  # should we use the (precomputed) approximation error model
    'useSpatialCorrelations': True,  # should we use the spatial correlation models for aerosol properties
    'quantifyUncertainty': False,  # should we quantify the uncertainties, BAR can provide pixel-level uncertainty estimates but by default these are not saved to save time and disc space
}

AODprior = {
    'range': 50.0,  # correlation length (km)
    'p': 1.5,  # p parameter (defines smoothness of the field)
    'sill': 0.10,  # variance for the spatially correlated component
    'nugget': 0.0025,  # local variance
}

FMFprior = {
    'range': 50.0,  # correlation length (km)
    'p': 1.5,  # p parameter (defines smoothness of the field)
    'sill': 0.25,  # variance for the spatially correlated component
    'nugget': 0.0100,  # local variance
}

surfacePriorModifier = {
    'meancoef': 1.0,  # surface prior model expected values are multiplied by this number, > 1.0 brighten surface reflectance assumptions, < 1.0 darken surface reflectance assumptions
    'stdcoef': 1.0,  # surface prior model standard deviations are multiplied by this number, > 1.0 adds uncertainty to surface prior, < 1.0 reduces uncertainty
}

# parse command line options
for option in sys.argv:
    if option[0] != '-':
        continue

    if option.lower() == '-quantifyuncertainty':
        settings['quantifyUncertainty'] = True
        print('  Uncertainty quantification enabled')
    elif option.lower() == '-nospatialcorrelation':
        settings['useSpatialCorrelations'] = False
        print('  Spatial correlation models disabled')
    elif option.lower() == '-noapproximationerror':
        settings['useApproximationErrorModel'] = False
        print('  Approximation error models disabled')
    elif option.lower()[:10] == '-aodprior=':
        print('  Using custom AOD prior parameters')
        try:
            l = json.loads(option[10:])
            for k, v in l.items():
                print('    {}: {}'.format(k, v))
                if k not in AODprior:
                    print('\n{} is not AOD prior parameter\n'.format(k))
                    sys.exit(1)
                AODprior[k] = v
        except:
            print('\nAOD prior parameter JSON is not valid.\n')
            sys.exit(1)
    elif option.lower()[:10] == '-fmfprior=':
        print('  Using custom FMF prior parameters')
        try:
            l = json.loads(option[10:])
            for k, v in l.items():
                print('    {}: {}'.format(k, v))
                if k not in FMFprior:
                    print('\n{} is not FMF prior parameter\n'.format(k))
                    sys.exit(1)
                FMFprior[k] = v
        except:
            print('\nFMF prior parameter JSON is not valid.\n')
            sys.exit(1)
    elif option.lower()[:14] == '-surfaceprior=':
        print('  Modifying surface reflectance prior')
        try:
            l = json.loads(option[14:])
            for k, v in l.items():
                print('    {}: {}'.format(k, v))
                if k not in surfacePriorModifier:
                    print('\n{} is not surface reflectance prior parameter\n'.format(k))
                    sys.exit(1)
                surfacePriorModifier[k] = v
        except:
            print('\nSurface reflectance prior modifier JSON is not valid.\n')
            sys.exit(1)
    else:
        print('\nUnknown command line option "{}"\n'.format(option))
        printUse()
        sys.exit(1)

# load approximation error model
try:
    with open(os.path.join('models', 'ApproximationError', 'approximationerror.json'), 'rt') as f:
        approximationErrorModel = json.load(f)
except:
    print('\nApproximation error file not found. This data should be included with the codes and data. Please check everything is correct.\n')
    sys.exit(1)

# load FMF prior model expected value (the file is not provided with the code, please see the readme for information how to download the file)
try:
    FMFpriorFile = os.path.join('models', 'FMF', 'gt_ff_0550nm.nc')
    f = Dataset(FMFpriorFile, 'r', format='NETCDF4')
    fmfMonth, fmfLat, fmfLon, fmfValues = np.arange(1, 13), f['lat'][:], f['lon'][:], f['aer_data'][:]
    f.close()
except:
    print('\nFMF prior model data file ({}) not found. This is an external file that needs to be downloaded before the code can be run.\nYou can download the file from: ftp://ftp-projects.zmaw.de/aerocom/climatology/MACv2_2017/550nm_2005/gt_ff_0550nm.nc\nPlease see README.txt for more information.\n'.format(FMFpriorFile))
    sys.exit(1)

# load AOD prior model expected value (the file is not provided with the code, please see the readme for information how to download the file)
try:
    AODpriorFile = os.path.join('models', 'AOD', 'gt_t_00550nm.nc')
    f = Dataset(AODpriorFile, 'r', format='NETCDF4')
    aodMonth, aodLat, aodLon, aodValues = np.arange(1, 13), f['lat'][:], f['lon'][:], f['aod'][:]
    f.close()
except:
    print('\nAOD prior model data file ({}) not found. This is an external file that needs to be downloaded before the code can be run.\nYou can download the file from: ftp://ftp-projects.zmaw.de/aerocom/climatology/MACv2_2017/550nm_2005/gt_t_00550nm.nc\nPlease see README.txt for more information.\n'.format(AODpriorFile))
    sys.exit(1)

# construct an interpolator for the expected values
AODFMFMeanIntpolator = {
    'AOD': RegularGridInterpolator((aodMonth, aodLat[::-1], aodLon), aodValues[:, ::-1, :]),
    'FMF': RegularGridInterpolator((fmfMonth, fmfLat[::-1], fmfLon), fmfValues[:, ::-1, :])
}

# load MODIS granule file
try:
    modisData = loadGranule(dataFile)
except:
    print('\nError loading datafile "{}". The file should be a valid MODIS MOD04_L2 or MYD04_L2 product file.\n'.format(dataFile))
    sys.exit(1)

granuleMonth, granuleYear = modisData['scantime'].month, modisData['scantime'].year
print('  Granule date and time: {}'.format(modisData['scantime'].strftime('%d %b %Y %H:%MZ (doy: %j)')))

try:
    granuleCenterPoint = [np.mean(modisData['lon']), np.mean(modisData['lat'])]
    print('  Granule center point: LAT: {:.2f}  LON: {:.2f}'.format(granuleCenterPoint[1], granuleCenterPoint[0]))
    area = list(filter(lambda x: x[1]['min_lat'] <= granuleCenterPoint[1] and
                x[1]['max_lat'] >= granuleCenterPoint[1] and
                x[1]['min_lon'] <= granuleCenterPoint[0] and
                x[1]['max_lon'] >= granuleCenterPoint[0], areas.items()))[0][0]
    print('  Area to be used for retrieval: {}'.format(area))
except:
    print('\nERROR in finding the area.\n')
    sys.exit(1)

# load surface reflectance prior model
try:
    surfacePrior = loadSurfaceReflectancePrior(granuleMonth, granuleYear, surfacePriorModifier, settings)
except:
    print('\nError loading surface prior.\n')
    sys.exit(1)

# prepare radiative transfer models
polynomialLookups, granuleMASK_full, allpartsData, gFull, aerosolmap = prepareModels(modisData, settings)
if len(allpartsData) > 1:
    print('  Retrieval will be run in {} parts'.format(len(allpartsData)))

savedata = {
    'BAR_AOD': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD466': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD644': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_2.5': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_5.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_16.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_84.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_95.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_AOD_percentile_97.5': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_2.5': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_5.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_16.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_84.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_95.0': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_percentile_97.5': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos466': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos553': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos644': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos211': -9.999 * np.ones(gFull.shape[0]),
    'BAR_logAODplus1': -9.999 * np.ones(gFull.shape[0]),
    'BAR_logAODplus1_std': -9.999 * np.ones(gFull.shape[0]),
    'BAR_FMF_std': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos466_std': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos553_std': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos644_std': -9.999 * np.ones(gFull.shape[0]),
    'BAR_rhos211_std': -9.999 * np.ones(gFull.shape[0]),
    'lon': gFull[:, 0],
    'lat': gFull[:, 1],
    'DarkTarget_AOD': modisData['Corrected_Optical_Depth_Land553'],
    'DarkTarget_FMF': modisData['Optical_Depth_Ratio_Small_Land'],
}

if not settings['quantifyUncertainty']:
    savedata = {k: savedata[k] for k in savedata if k.find('percentile') < 0 and k.find('_std') < 0}

for iPart, partData in enumerate(allpartsData):
    if len(allpartsData) == 1:
        print('  Running retrieval...')
    else:
        print('  Running retrieval... (part {:d}/{:d})'.format(iPart + 1, len(allpartsData)))

    # preparing masks & coordinates
    MASKl, MASKf = partData['MASKl'], partData['MASKf']
    g = gFull[MASKf, :]

    # preparing surface reflectance prior for this granule
    indices = surfacePrior['kNNmean'].kneighbors(g, return_distance=False)
    indicesstd = surfacePrior['kNNstd'].kneighbors(g, return_distance=False)

    granuleSurfacePrior = {
        'mean466': surfacePrior['surfprior466'][indices].mean(axis=1),
        'mean550': surfacePrior['surfprior550'][indices].mean(axis=1),
        'mean644': surfacePrior['surfprior644'][indices].mean(axis=1),
        'mean211': surfacePrior['surfprior211'][indices].mean(axis=1),
        'std466': np.sqrt(np.clip(surfacePrior['surfprior466'][indices].std(axis=1), a_min=1.0e-6, a_max=1.0)**2 + np.mean(surfacePrior['surfprior466std'][indicesstd]**2, axis=1)),
        'std550': np.sqrt(np.clip(surfacePrior['surfprior550'][indices].std(axis=1), a_min=1.0e-6, a_max=1.0)**2 + np.mean(surfacePrior['surfprior550std'][indicesstd]**2, axis=1)),
        'std644': np.sqrt(np.clip(surfacePrior['surfprior644'][indices].std(axis=1), a_min=1.0e-6, a_max=1.0)**2 + np.mean(surfacePrior['surfprior644std'][indicesstd]**2, axis=1)),
        'std211': np.sqrt(np.clip(surfacePrior['surfprior211'][indices].std(axis=1), a_min=1.0e-6, a_max=1.0)**2 + np.mean(surfacePrior['surfprior211std'][indicesstd]**2, axis=1)),
    }
    Nvars = g.shape[0]

    # preparing observed data
    REFW466 = modisData['rho466'][MASKl]
    STDREFW466 = modisData['stdrho466'][MASKl]
    REFW553 = modisData['rho550'][MASKl]
    STDREFW553 = modisData['stdrho550'][MASKl]
    REFW644 = modisData['rho644'][MASKl]
    STDREFW644 = modisData['stdrho644'][MASKl]
    REFW2113 = modisData['rho211'][MASKl]
    STDREFW2113 = modisData['stdrho211'][MASKl]
    MeasData = np.concatenate((REFW466, REFW553, REFW644, REFW2113))
    MeasDataStd = np.concatenate((STDREFW466, STDREFW553, STDREFW644, STDREFW2113))

    # bounds for AOD & FMF (& surface reflectance)
    min_AOD, max_AOD = np.log(0.0 + 1.0), np.log(5.0 + 1.0)
    min_FMF, max_FMF = 0.0, 1.0
    bounds = np.array([[min_AOD, max_AOD], [min_FMF, max_FMF], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0], [0.0, 1.0]])

    # load approximation error model
    if settings['useApproximationErrorModel']:
        print('    Using approximation error model {}/month {}'.format(area, granuleMonth))
        MeasNoiseVar = (1.0 / (1.0 + MeasData) * MeasDataStd)**2  # LOGSCALE RHOS
        Nwavelengths = 4
        icov_ii, icov_jj, icov_vv = np.zeros(Nwavelengths**2 * Nvars, dtype=int), np.zeros(Nwavelengths**2 * Nvars, dtype=int), np.zeros(Nwavelengths**2 * Nvars)
        for ii in range(Nvars):
            local_cov = np.array(np.array(approximationErrorModel[area][str(granuleMonth)]['cov']).reshape((Nwavelengths, Nwavelengths)) + diags([MeasNoiseVar[ii], MeasNoiseVar[ii + 1 * Nvars], MeasNoiseVar[ii + 2 * Nvars], MeasNoiseVar[ii + 3 * Nvars]], 0))
            local_icov = np.linalg.inv(local_cov)
            icov_ii[ii * Nwavelengths**2:(ii + 1) * Nwavelengths**2] = np.concatenate([[ii + 0 * Nvars] * Nwavelengths, [ii + 1 * Nvars] * Nwavelengths, [ii + 2 * Nvars] * Nwavelengths, [ii + 3 * Nvars] * Nwavelengths])
            icov_jj[ii * Nwavelengths**2:(ii + 1) * Nwavelengths**2] = np.concatenate([[ii + 0 * Nvars, ii + 1 * Nvars, ii + 2 * Nvars, ii + 3 * Nvars] * Nwavelengths])
            icov_vv[ii * Nwavelengths**2:(ii + 1) * Nwavelengths**2] = local_icov.ravel()

        MeasNoiseiCov = csr_matrix((icov_vv, (icov_ii, icov_jj)), shape=(Nwavelengths * Nvars, Nwavelengths * Nvars))
        MeasNoiseE = np.tile(np.array(approximationErrorModel[area][str(granuleMonth)]['E'])[:, np.newaxis], (1, Nvars)).ravel()
    else:
        MeasNoiseCov = diags((1.0 / (1.0 + MeasData) * MeasDataStd)**2, 0)  # LOGSCALE RHOS
        MeasNoiseiCov = diags(1.0 / MeasNoiseCov.diagonal(), 0)
        MeasNoiseE = 0.0
    # switch to log-scale measurement data (TOA reflectances)
    MeasData = np.log(MeasData + 1.0)

    # Initial values for retrieval
    AOD0 = AODFMFMeanIntpolator['AOD'](np.array([[granuleMonth, gg[1], gg[0]] for gg in np.clip(g, a_min=[-179.5, -89.5], a_max=[179.5, 89.5])]))
    FMF0 = AODFMFMeanIntpolator['FMF'](np.array([[granuleMonth, gg[1], gg[0]] for gg in np.clip(g, a_min=[-179.5, -89.5], a_max=[179.5, 89.5])]))

    # prepare the lookup-tables to final form
    polynomiallookupMatrices_fine = combinePolyLookupList(polynomialLookups[MASKl])
    polynomiallookupMatrices_coarse = combinePolyLookupList(polynomialLookups[MASKl], 4)
    BARresults = BARretrieve(g, AOD0, FMF0, AODprior, FMFprior, granuleSurfacePrior,
                             polynomiallookupMatrices_fine,
                             polynomiallookupMatrices_coarse,
                             MeasData, MeasNoiseiCov, MeasNoiseE, MASKf, bounds, settings)

    for k, v in BARresults.items():
        savedata[k][MASKf] = v

outputfile = os.path.join(outputFilePath, outputFileName)
print('  Saving results to {}'.format(outputfile))
nc = Dataset(outputfile, 'w')
nc.AODprior = str(AODprior).replace('\'', '')
nc.FMFprior = str(FMFprior).replace('\'', '')
nc.surfacePriorModifier = str(surfacePriorModifier).replace('\'', '')
nc.retrievalsettings = str(settings).replace('\'', '')
nc.inputfilename = inputFileName
nc.granuleTime = modisData['scantime'].strftime('%d-%m-%Y %H:%MZ')
nc.retrievalTime = datetime.utcnow().strftime('%d-%m-%Y %H:%MZ')
nX, nY = modisData['nX'], modisData['nY']
ncX = nc.createDimension('X', nX)
ncY = nc.createDimension('Y', nY)
var = nc.createVariable('latitude', 'f4', ('X', 'Y'))
var[:] = savedata['lat'].reshape((nX, nY))
var = nc.createVariable('longitude', 'f4', ('X', 'Y'))
var[:] = savedata['lon'].reshape((nX, nY))
for key in savedata.keys():
    if not (key[:3] == 'BAR' or key[:10] == 'DarkTarget'):
        continue
    var = nc.createVariable(key, 'f4', ('X', 'Y'))
    varToBeSaved = savedata[key]
    if key.find('AOD') >= 0:
        var.vmin = -0.1
        var.vmax = 5.0
    else:
        var.vmin = 0.0
        var.vmax = 1.0
    varToBeSaved[np.logical_or(varToBeSaved < var.vmin, varToBeSaved > var.vmax)] = np.nan
    var[:] = varToBeSaved.reshape((nX, nY))
nc.close()
print('***** ALL DONE FOR THIS GRANULE *****\n\nThanks for using Bayesian Aerosol Retrieval!\n')

ext: py

max_stars_repo_name: TUT-ISI/BARalgorithm

max_issues_repo_name: TUT-ISI/BARalgorithm

max_forks_repo_name: TUT-ISI/BARalgorithm

reason: import numpy,from scipy


text:
import os
import glob
import time
import torch
import torch.utils
import torch.nn as nn
import torchvision
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.nn.modules.module import _addindent
import numpy as np
from collections import OrderedDict
import argparse

from dataset import  GRID_1D_lstm_landmark_pca, LRW_1D_lstm_landmark_pca, GRID_1D_single_landmark_pca, LRW_1D_single_landmark_pca

from models import AT_single, AT_net

from torch.nn import init
import utils
def multi2single(model_path, id):
    checkpoint = torch.load(model_path)
    state_dict = checkpoint['state_dict']
    if id ==1:
        new_state_dict = OrderedDict()
        for k, v in state_dict.items():
            name = k[7:]
            new_state_dict[name] = v
        return new_state_dict
    else:
        return state_dict

def initialize_weights( net, init_type='normal', gain=0.02):
        def init_func(m):
            classname = m.__class__.__name__
            if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
                if init_type == 'normal':
                    init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'kaiming':
                    init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    init.orthogonal_(m.weight.data, gain=gain)
                else:
                    raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
                if hasattr(m, 'bias') and m.bias is not None:
                    init.constant_(m.bias.data, 0.0)
            elif classname.find('BatchNorm2d') != -1:
                init.normal_(m.weight.data, 1.0, gain)
                init.constant_(m.bias.data, 0.0)

        print('initialize network with %s' % init_type)
        net.apply(init_func)

class Trainer():
    def __init__(self, config):
        if config.lstm == True:
            self.generator = AT_net()
        else:
            self.generator = AT_single()
        self.l1_loss_fn =  nn.L1Loss()
        self.mse_loss_fn = nn.MSELoss()
        self.config = config

        if config.cuda:
            device_ids = [int(i) for i in config.device_ids.split(',')]
            self.generator     = nn.DataParallel(self.generator, device_ids=device_ids).cuda()
            # self.generator     = self.generator.cuda()
            self.mse_loss_fn   = self.mse_loss_fn.cuda()
            self.l1_loss_fn = self.l1_loss_fn.cuda()
# #########single GPU#######################

#         if config.cuda:
#             device_ids = [int(i) for i in config.device_ids.split(',')]
#             self.generator     = self.generator.cuda()
#             self.encoder = self.encoder.cuda()
#             self.mse_loss_fn   = self.mse_loss_fn.cuda()
#             self.l1_loss_fn =  nn.L1Loss().cuda()
        initialize_weights(self.generator)
        self.start_epoch = 0
        if config.load_model:
            self.start_epoch = config.start_epoch
            self.load(config.pretrained_dir, config.pretrained_epoch)
        self.opt_g = torch.optim.Adam( self.generator.parameters(),
            lr=config.lr, betas=(config.beta1, config.beta2))
        if config.lstm:
            if config.dataset == 'lrw':
                self.dataset = LRW_1D_lstm_landmark_pca(config.dataset_dir, train=config.is_train)
            else:
                self.dataset = GRID_1D_lstm_landmark_pca(config.dataset_dir, train=config.is_train)
        else:
            if config.dataset == 'lrw':
                self.dataset = LRW_1D_single_landmark_pca(config.dataset_dir, train=config.is_train)
            else:
                self.dataset = GRID_1D_single_landmark_pca(config.dataset_dir, train=config.is_train)

        # elif config.dataset == 'ldc':
        #     self.dataset = LDCDataset(config.dataset_dir, train=config.is_train)

        self.data_loader = DataLoader(self.dataset,
                                      batch_size=config.batch_size,
                                      num_workers=config.num_thread,
                                      shuffle=True, drop_last=True)

    def fit(self):
        config = self.config
        if config.dataset == 'grid':

            pca = torch.FloatTensor( np.load('../basics/U_grid.npy')[:,:6]).cuda()
            mean =torch.FloatTensor( np.load('../basics/mean_grid.npy')).cuda()
        elif config.dataset == 'lrw':

            pca = torch.FloatTensor(np.load('../basics/U_lrw1.npy')[:,:6] ).cuda()
            mean = torch.FloatTensor(np.load('../basics/mean_lrw1.npy')).cuda()
        else:
            raise Exception('wrong key word for the dataset input')

        num_steps_per_epoch = len(self.data_loader)
        cc = 0
        t0 = time.time()
        xLim=(-1.0, 1.0)
        yLim=(-1.0, 1.0)
        xLab = 'x'
        yLab = 'y'

        for epoch in range(self.start_epoch, config.max_epochs):
            for step, (example_landmark, example_audio, lmark, audio) in enumerate(self.data_loader):
                t1 = time.time()

                if config.cuda:
                    lmark    = Variable(lmark.float()).cuda()
                    audio = Variable(audio.float()).cuda()
                    example_landmark = Variable(example_landmark.float()).cuda()
                else:
                    lmark    = Variable(lmark.float())
                    audio = Variable(audio.float())
                    example_landmark = Variable(example_landmark.float())

                fake_lmark= self.generator( example_landmark, audio)

                loss =  self.mse_loss_fn(fake_lmark , lmark)
                loss.backward()
                self.opt_g.step()
                self._reset_gradients()

                if (step+1) % 10 == 0 or (step+1) == num_steps_per_epoch:
                    steps_remain = num_steps_per_epoch-step+1 + \
                        (config.max_epochs-epoch+1)*num_steps_per_epoch

                    print("[{}/{}][{}/{}]   loss1: {:.8f},data time: {:.4f},  model time: {} second"
                          .format(epoch+1, config.max_epochs,
                                  step+1, num_steps_per_epoch, loss,  t1-t0,  time.time() - t1))
                if (step) % (int(num_steps_per_epoch  / 10 )) == 0 and step != 0:
                    if config.lstm:
                        lmark = lmark.view(lmark.size(0) * lmark.size(1) , 6)
                        lmark = torch.mm( lmark, pca.t() )
                        lmark = lmark + mean.expand_as(lmark)
                        fake_lmark = fake_lmark.view(fake_lmark.size(0) * fake_lmark.size(1) , 6)
                        fake_lmark[:, 1:6] *= 2*torch.FloatTensor(np.array([1.1, 1.2, 1.3, 1.4, 1.5])).cuda()
                        fake_lmark = torch.mm( fake_lmark, pca.t() )
                        fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
                        lmark = lmark.view(config.batch_size, 16, 136)
                        lmark = lmark.data.cpu().numpy()
                        fake_lmark = fake_lmark.view(config.batch_size, 16, 136)
                        fake_lmark = fake_lmark.data.cpu().numpy()
                        for indx in range(3):
                            for jj in range(16):

                                name = "{}real_{}_{}_{}.png".format(config.sample_dir,cc, indx,jj)
                                utils.plot_flmarks(lmark[indx,jj], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
                                name = "{}fake_{}_{}_{}.png".format(config.sample_dir,cc, indx,jj)
                                utils.plot_flmarks(fake_lmark[indx,jj], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
                        torch.save(self.generator.state_dict(),
                                   "{}/atnet_lstm_{}.pth"
                                   .format(config.model_dir,cc))
                    else:
                        lmark = torch.mm( lmark, pca.t() )
                        lmark = lmark + mean.expand_as(lmark)
                        fake_lmark[ 1:6] *= 2*torch.FloatTensor(np.array([1.1, 1.2, 1.3, 1.4, 1.5])).cuda()
                        fake_lmark = torch.mm( fake_lmark, pca.t() )
                        fake_lmark = fake_lmark + mean.expand_as(fake_lmark)
                        lmark = lmark.data.cpu().numpy()
                        fake_lmark = fake_lmark.data.cpu().numpy()
                        for indx in range(0,config.batch_size, 8):
                            name = "{}single_real_{}_{}.png".format(config.sample_dir,cc, indx)
                            utils.plot_flmarks(lmark[indx], name, xLim, yLim, xLab, yLab, figsize=(10, 10))
                            name = "{}single_fake_{}_{}.png".format(config.sample_dir,cc, indx)
                            utils.plot_flmarks(fake_lmark[indx], name, xLim, yLim, xLab, yLab, figsize=(10, 10))

                        torch.save(self.generator.state_dict(),
                                   "{}anet_single_{}.pth"
                                   .format(config.model_dir,cc))
                    cc += 1

                t0 = time.time()

    def _reset_gradients(self):
        self.generator.zero_grad()

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--lr",
                        type=float,
                        default=0.0002)
    parser.add_argument("--beta1",
                        type=float,
                        default=0.5)
    parser.add_argument("--beta2",
                        type=float,
                        default=0.999)
    parser.add_argument("--lambda1",
                        type=int,
                        default=100)
    parser.add_argument("--batch_size",
                        type=int,
                        default=16)
    parser.add_argument("--max_epochs",
                        type=int,
                        default=10)
    parser.add_argument("--cuda",
                        default=True)
    parser.add_argument("--dataset_dir",
                        type=str,
                        default="../dataset/")
                        # default="/mnt/ssd0/dat/lchen63/grid/pickle/")
                        # default = '/media/lele/DATA/lrw/data2/pickle')
    parser.add_argument("--model_dir",
                        type=str,
                        default="../model/atnet/")
                        # default="/mnt/disk1/dat/lchen63/grid/model/model_gan_r")
                        # default='/media/lele/DATA/lrw/data2/model')
    parser.add_argument("--sample_dir",
                        type=str,
                        default="../sample/atnet/")
                        # default="/mnt/disk1/dat/lchen63/grid/sample/model_gan_r/")
                        # default='/media/lele/DATA/lrw/data2/sample/lstm_gan')
    parser.add_argument('--device_ids', type=str, default='0')
    parser.add_argument('--dataset', type=str, default='lrw')
    parser.add_argument('--lstm', type=bool, default= True)
    parser.add_argument('--num_thread', type=int, default=2)
    parser.add_argument('--weight_decay', type=float, default=4e-4)
    parser.add_argument('--load_model', action='store_true')
    parser.add_argument('--pretrained_dir', type=str)
    parser.add_argument('--pretrained_epoch', type=int)
    parser.add_argument('--start_epoch', type=int, default=0, help='start from 0')
    parser.add_argument('--rnn', type=bool, default=True)

    return parser.parse_args()

def main(config):
    t = trainer.Trainer(config)
    t.fit()

if __name__ == "__main__":

    config = parse_args()
    # os.environ["CUDA_VISIBLE_DEVICES"] = config.device_ids
    config.is_train = 'train'
    import atnet as trainer
    if not os.path.exists(config.model_dir):
        os.mkdir(config.model_dir)
    if not os.path.exists(config.sample_dir):
        os.mkdir(config.sample_dir)
    config.cuda1 = torch.device('cuda:{}'.format(config.device_ids))
    main(config)

max_stars_repo_name: Gwiths/ATVGnet

max_issues_repo_name: Gwiths/ATVGnet

max_forks_repo_name: Gwiths/ATVGnet

reason: import numpy


text:
# Copyright 2021 Fagner Cunha
#
# 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.

import json
import os
import random

from absl import app
from absl import flags
import numpy as np
import tensorflow as tf

from iwildcamlib import CategoryMap, generate_submission
import bags
import dataloader
import model_builder

os.environ['TF_DETERMINISTIC_OPS'] = '1'

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'model_name', default='efficientnet-b0',
    help=('Model name of the archtecture'))

flags.DEFINE_integer(
    'input_size', default=224,
    help=('Input size of the model'))

flags.DEFINE_bool(
    'use_bags', default=False,
    help=('Use Balanced Group Softmax to train model'))

flags.DEFINE_bool(
    'use_flip_image', default=False,
    help=('Add flipped images to predictions'))

flags.DEFINE_integer(
    'empty_class_id', default=0,
    help=('Empty class id for balanced group softmax'))

flags.DEFINE_integer(
    'batch_size', default=32,
    help=('Batch size used during training.'))

flags.DEFINE_string(
    'ckpt_dir_full', default=None,
    help=('Location of the checkpoint files for the model trained on full'
          ' images'))

flags.DEFINE_string(
    'ckpt_dir_bbox', default=None,
    help=('Location of the checkpoint files for the model trained on bbox'
          ' crops'))

flags.DEFINE_float(
    'bbox_model_weight', default=0.7,
    help=('Weight given to bounding box model on combined prediction. Full'
          ' image will have weight of 1 - bbox_model_weight'))

flags.DEFINE_string(
    'annotations_json', default=None,
    help=('Path to json file containing the training annotations json for'
          ' the iWildCam2021 competition'))

flags.DEFINE_string(
    'train_dataset_split', default=None,
    help=('Path to json file containing the train/validation split based on'
          ' locations.'))

flags.DEFINE_string(
    'test_info_json', default=None,
    help=('Path to json file containing the test information json for'
          ' the iWildCam2021 competition'))

flags.DEFINE_string(
    'dataset_dir', default=None,
    help=('Path to directory containing test images.'))

flags.DEFINE_string(
    'megadetector_results_json', default=None,
    help=('Path to json file containing megadetector results.'))

flags.DEFINE_string(
    'submission_file_path', default=None,
    help=('File name to save predictions on iWildCam2020 results format.'))

flags.DEFINE_integer(
    'log_frequence', default=500,
    help=('Log prediction every n steps'))

if 'random_seed' not in list(FLAGS):
  flags.DEFINE_integer(
      'random_seed', default=42,
      help=('Random seed for reproductible experiments'))

flags.mark_flag_as_required('ckpt_dir_full')
flags.mark_flag_as_required('ckpt_dir_bbox')
flags.mark_flag_as_required('annotations_json')
flags.mark_flag_as_required('test_info_json')
flags.mark_flag_as_required('dataset_dir')
flags.mark_flag_as_required('megadetector_results_json')
flags.mark_flag_as_required('submission_file_path')

def load_train_validation_split():
  if FLAGS.train_dataset_split is None:
    return None, None

  with tf.io.gfile.GFile(FLAGS.train_dataset_split, 'r') as json_file:
    json_data = json.load(json_file)

  return json_data['train'], json_data['validation']

def _load_model(num_classes, bal_group_softmax=None):
  model_bbox = model_builder.create(model_name=FLAGS.model_name,
                              num_classes=num_classes,
                              input_size=FLAGS.input_size,
                              unfreeze_layers=0,
                              bags=bal_group_softmax)
  checkpoint_path = os.path.join(FLAGS.ckpt_dir_bbox, "ckp")
  model_bbox.load_weights(checkpoint_path)
  if bal_group_softmax is not None:
    model_bbox = bal_group_softmax.create_prediction_model(model_bbox)

  model_full = model_builder.create(model_name=FLAGS.model_name,
                              num_classes=num_classes,
                              input_size=FLAGS.input_size,
                              unfreeze_layers=0,
                              bags=bal_group_softmax)
  checkpoint_path = os.path.join(FLAGS.ckpt_dir_full, "ckp")
  model_full.load_weights(checkpoint_path)
  if bal_group_softmax is not None:
    model_full = bal_group_softmax.create_prediction_model(model_full)

  bbox_weight = FLAGS.bbox_model_weight
  full_weight = 1 - FLAGS.bbox_model_weight
  if FLAGS.use_flip_image:
    bbox_weight /= 2.0
    full_weight /= 2.0

  image_bbox = tf.keras.Input(shape=(FLAGS.input_size, FLAGS.input_size, 3))
  x1 = model_bbox(image_bbox, training=False)
  x1 = tf.keras.layers.Lambda(lambda x: x * bbox_weight)(x1)

  image_full = tf.keras.Input(shape=(FLAGS.input_size, FLAGS.input_size, 3))
  x2 = model_full(image_full, training=False)
  x2 = tf.keras.layers.Lambda(lambda x: x * full_weight)(x2)

  if FLAGS.use_flip_image:
    image_bboxf = tf.keras.Input(shape=(FLAGS.input_size, FLAGS.input_size, 3))
    x1f = model_bbox(image_bboxf, training=False)
    x1f = tf.keras.layers.Lambda(lambda x: x * bbox_weight)(x1f)

    image_fullf = tf.keras.Input(shape=(FLAGS.input_size, FLAGS.input_size, 3))
    x2f = model_full(image_fullf, training=False)
    x2f = tf.keras.layers.Lambda(lambda x: x * full_weight)(x2f)

    inputs = [image_bbox, image_bboxf, image_full, image_fullf]
    outputs = tf.keras.layers.Add()([x1, x1f, x2, x2f])
  else:
    inputs = [image_bbox, image_full]
    outputs = tf.keras.layers.Add()([x1, x2])
  model = tf.keras.models.Model(inputs=inputs, outputs=[outputs])

  return model

def _build_input_data(category_map):
  input_data =  dataloader.JsonWBBoxInputProcessor(
    dataset_json=FLAGS.test_info_json,
    dataset_dir=FLAGS.dataset_dir,
    megadetector_results_json=FLAGS.megadetector_results_json,
    batch_size=FLAGS.batch_size,
    batch_drop_remainder=False,
    category_map=category_map,
    is_training=False,
    output_size=FLAGS.input_size,
    crop_mode='both_flip' if FLAGS.use_flip_image else 'both',
    provide_instance_id=True,
    seed=FLAGS.random_seed)

  return input_data.make_source_dataset()

def predict_classifier(model, dataset):
  instance_ids = []
  predictions = []
  count = 0

  for batch, metadata in dataset:
    pred = model(batch, training=False)
    _, instanceid = metadata
    instance_ids += list(instanceid.numpy())
    predictions += list(pred.numpy())

    if count % FLAGS.log_frequence == 0:
      tf.compat.v1.logging.info('Finished eval step %d' % count)
    count += 1

  return instance_ids, predictions

def format_instance_id(instance_ids):
  instance_ids = [filename.decode("utf-8")[:-4] for filename in instance_ids]

  return instance_ids

def set_random_seeds():
  random.seed(FLAGS.random_seed)
  np.random.seed(FLAGS.random_seed)
  tf.random.set_seed(FLAGS.random_seed)

def main(_):
  set_random_seeds()

  category_map = CategoryMap(FLAGS.annotations_json)
  train_loc, _ = load_train_validation_split()
  bal_group_softmax = bags.BalancedGroupSoftmax(
        FLAGS.annotations_json,
        category_map,
        FLAGS.empty_class_id,
        selected_locations=train_loc) if FLAGS.use_bags else None
  dataset, _ = _build_input_data(category_map)
  model = _load_model(category_map.get_num_classes(), bal_group_softmax)

  instance_ids, predictions = predict_classifier(model, dataset)
  instance_ids = format_instance_id(instance_ids)

  generate_submission(instance_ids, predictions, category_map,
                      FLAGS.test_info_json, FLAGS.megadetector_results_json,
                      FLAGS.submission_file_path)

if __name__ == '__main__':
  app.run(main)


text:
import numpy as np

from .. import files
from ..validators import column_or_1d

from . import functional as F

__all__ = [
    "AddChannel",
    "PadToWindow",
    "StandardizeSpect",
    "ToFloatTensor",
    "ToLongTensor",
    "ViewAsWindowBatch",
]

# adapted from:
# https://github.com/NickleDave/hybrid-vocal-classifier/blob/master/hvc/neuralnet/utils.py
class StandardizeSpect:
    """transform that standardizes spectrograms so they are all
    on the same scale, by subtracting off the mean and dividing by the
    standard deviation from a 'fit' set of spectrograms.

    Attributes
    ----------
    mean_freqs : numpy.ndarray
        mean values for each frequency bin across the fit set of spectrograms
    std_freqs : numpy.ndarray
        standard deviation for each frequency bin across the fit set of spectrograms
    non_zero_std : numpy.ndarray
        boolean, indicates where std_freqs has non-zero values. Used to avoid divide-by-zero errors.
    """

    def __init__(self, mean_freqs=None, std_freqs=None, non_zero_std=None):
        """initialize a new StandardizeSpect instance

        Parameters
        ----------
        mean_freqs : numpy.ndarray
            vector of mean values for each frequency bin across the fit set of spectrograms
        std_freqs : numpy.ndarray
            vector of standard deviations for each frequency bin across the fit set of spectrograms
        non_zero_std : numpy.ndarray
            boolean, indicates where std_freqs has non-zero values. Used to avoid divide-by-zero errors.
        """
        if any([arg is not None for arg in (mean_freqs, std_freqs, non_zero_std)]):
            mean_freqs, std_freqs, non_zero_std = (
                column_or_1d(arr) for arr in (mean_freqs, std_freqs, non_zero_std)
            )
            if (
                len(
                    np.unique(
                        [arg.shape[0] for arg in (mean_freqs, std_freqs, non_zero_std)]
                    )
                )
                != 1
            ):
                raise ValueError(
                    f"mean_freqs, std_freqs, and non_zero_std must all have the same length"
                )

        self.mean_freqs = mean_freqs
        self.std_freqs = std_freqs
        self.non_zero_std = non_zero_std

    @classmethod
    def fit_df(cls, df, spect_key="s"):
        """fits StandardizeSpect instance, given a pandas DataFrame representing a dataset

        Parameters
        ----------
        df : pandas.DataFrame
            loaded from a .csv file representing a dataset, created by vak.io.dataframe.from_files
        spect_key : str
            key in files in 'spect_path' column that maps to spectrograms in arrays.
            Default is 's'.

        Returns
        -------
        standardize_spect : StandardizeSpect
            instance fit to spectrograms in df
        """
        spect_paths = df["spect_path"]
        spect = files.spect.load(spect_paths[0])[spect_key]
        # in files, spectrograms are in orientation (freq bins, time bins)
        # so we take mean and std across columns, i.e. time bins, i.e. axis 1
        mean_freqs = np.mean(spect, axis=1)
        std_freqs = np.std(spect, axis=1)

        for spect_path in spect_paths[1:]:
            spect = files.spect.load(spect_path)[spect_key]
            mean_freqs += np.mean(spect, axis=1)
            std_freqs += np.std(spect, axis=1)
        mean_freqs = mean_freqs / len(spect_paths)
        std_freqs = std_freqs / len(spect_paths)
        non_zero_std = np.argwhere(std_freqs != 0)
        return cls(mean_freqs, std_freqs, non_zero_std)

    @classmethod
    def fit(cls, spect):
        """fit a StandardizeSpect instance.

        Parameters
        ----------
        spect : numpy.ndarray
            with dimensions (frequency bins, time bins)

        Notes
        -----
        Input should be spectrogram.
        Fit function finds the mean and standard deviation of each frequency bin,
        which are used by `transform` method to scale other spectrograms.
        """
        # TODO: make this function accept list and/or ndarray with batch dimension
        if spect.ndim != 2:
            raise ValueError("input spectrogram should be a 2-d array")

        mean_freqs = np.mean(spect, axis=1)
        std_freqs = np.std(spect, axis=1)
        assert mean_freqs.shape[-1] == spect.shape[0]
        assert std_freqs.shape[-1] == spect.shape[0]
        non_zero_std = np.argwhere(std_freqs != 0)
        return cls(mean_freqs, std_freqs, non_zero_std)

    def __call__(self, spect):
        """normalizes input spectrogram with fit parameters.

        Parameters
        ----------
        spect : numpy.ndarray
            2-d array with dimensions (frequency bins, time bins).

        Returns
        -------
        z_norm_spect : numpy.ndarray
            array standardized to same scale as set of spectrograms that
            SpectScaler was fit with
        """
        if any([not hasattr(self, attr) for attr in ["mean_freqs", "std_freqs"]]):
            raise AttributeError(
                "SpectScaler properties are set to None,"
                "must call fit method first to set the"
                "value of these properties before calling"
                "transform"
            )

        if type(spect) != np.ndarray:
            raise TypeError(
                f"type of spect must be numpy.ndarray but was: {type(spect)}"
            )

        if spect.shape[0] != self.mean_freqs.shape[0]:
            raise ValueError(
                f"number of rows in spects, {spect.shape[0]}, "
                f"does not match number of elements in self.mean_freqs, {self.mean_freqs.shape[0]},"
                "i.e. the number of frequency bins from the spectrogram"
                "to which the scaler was fit originally"
            )

        return F.standardize_spect(
            spect, self.mean_freqs, self.std_freqs, self.non_zero_std
        )

    def __repr__(self):
        args = f"(mean_freqs={self.mean_freqs}, std_freqs={self.std_freqs}, non_zero_std={self.non_zero_std})"
        return self.__class__.__name__ + args

class PadToWindow:
    """pad a 1d or 2d array so that it can be reshaped
    into consecutive windows of specified size

    Parameters
    ----------
    arr : numpy.ndarray
        with 1 or 2 dimensions, e.g. a vector of labeled timebins
        or a spectrogram.
    window_size : int
        width of window in number of elements.
    padval : float
        value to pad with. Added to end of array, the
        "right side" if 2-dimensional.
    return_padding_mask : bool
        if True, return a boolean vector to use for cropping
        back down to size before padding. padding_mask has size
        equal to width of padded array, i.e. original size
        plus padding at the end, and has values of 1 where
        columns in padded are from the original array,
        and values of 0 where columns were added for padding.

    Returns
    -------
    padded : numpy.ndarray
        padded with padval
    padding_mask : np.bool
        has size equal to width of padded, i.e. original size
        plus padding at the end. Has values of 1 where
        columns in padded are from the original array,
        and values of 0 where columns were added for padding.
        Only returned if return_padding_mask is True.
    """

    def __init__(self, window_size, padval=0.0, return_padding_mask=True):
        if not (type(window_size) == int) or (
            type(window_size) == float and window_size.is_integer() is False
        ):
            raise ValueError(
                f"window size must be an int or a whole number float;"
                f" type was {type(window_size)} and value was {window_size}"
            )

        if type(padval) not in (int, float):
            raise TypeError(
                f"type for padval must be int or float but was: {type(padval)}"
            )
        if not type(return_padding_mask) == bool:
            raise TypeError(
                "return_padding_mask must be boolean (True or False), "
                f"but was type {type(return_padding_mask)} with value {return_padding_mask}"
            )

        self.window_size = window_size
        self.padval = padval
        self.return_padding_mask = return_padding_mask

    def __call__(self, arr):
        return F.pad_to_window(
            arr, self.window_size, self.padval, self.return_padding_mask
        )

    def __repr__(self):
        args = f"(window_size={self.window_size}, padval={self.padval}, return_padding_mask={self.return_padding_mask})"
        return self.__class__.__name__ + args

class ViewAsWindowBatch:
    """return view of a 1d or 2d array as a batch of non-overlapping windows

    Parameters
    ----------
    arr : numpy.ndarray
        with 1 or 2 dimensions, e.g. a vector of labeled timebins
        or a 2-d array representing a spectrogram.
        If the array has 2-d dimensions, the returned array will
        have dimensions (batch, height of array, window width)
    window_width : int
        width of window in number of elements.

    Returns
    -------
    batch_windows : numpy.ndarray
        with shape (batch size, window_width) if array is 1d,
        or with shape (batch size, height, window_width) if array is 2d.
        Batch size will be arr.shape[-1] // window_width.
        Window width must divide arr.shape[-1] evenly.
        To pad the array so it can be divided into windows of the specified
        width, use the `pad_to_window` transform

    Notes
    -----
    adapted from skimage.util.view_as_blocks
    https://github.com/scikit-image/scikit-image/blob/f1b7cf60fb80822849129cb76269b75b8ef18db1/skimage/util/shape.py#L9
    """

    def __init__(self, window_width):
        if not (type(window_width) == int) or (
            type(window_width) == float and window_width.is_integer() is False
        ):
            raise ValueError(
                f"window size must be an int or a whole number float;"
                f" type was {type(window_width)} and value was {window_width}"
            )

        self.window_width = window_width

    def __call__(self, arr):
        return F.view_as_window_batch(arr, self.window_width)

    def __repr__(self):
        args = f"(window_width={self.window_width})"
        return self.__class__.__name__ + args

class ToFloatTensor:
    """convert Numpy array to torch.FloatTensor.

    Parameters
    ----------
    arr : numpy.ndarray

    Returns
    -------
    float_tensor
        with dtype 'float32'
    """

    def __init__(self):
        pass

    def __call__(self, arr):
        return F.to_floattensor(arr)

    def __repr__(self):
        return self.__class__.__name__

class ToLongTensor:
    """convert Numpy array to torch.LongTensor.

    Parameters
    ----------
    arr : numpy.ndarray

    Returns
    -------
    long_tensor : torch.Tensor
        with dtype 'float64'
    """

    def __init__(self):
        pass

    def __call__(self, arr):
        return F.to_longtensor(arr)

    def __repr__(self):
        return self.__class__.__name__

class AddChannel:
    """add a channel dimension to a 2-dimensional tensor.
    Transform that makes it easy to treat a spectrogram as an image,
    by adding a dimension with a single 'channel', analogous to grayscale.
    In this way the tensor can be fed to e.g. convolutional layers.

    Parameters
    ----------
    input : torch.Tensor
        with two dimensions (height, width).
    channel_dim : int
        dimension where "channel" is added.
        Default is 0, which returns a tensor with dimensions (channel, height, width).
    """

    def __init__(self, channel_dim=0):
        if not (type(channel_dim) == int) or (
            type(channel_dim) == float and channel_dim.is_integer() is False
        ):
            raise ValueError(
                f"window size must be an int or a whole number float;"
                f" type was {type(channel_dim)} and value was {channel_dim}"
            )

        channel_dim = int(channel_dim)

        if channel_dim < 0 and channel_dim != -1:
            raise ValueError(
                "value of channel_dim should be a non-negative integer, or -1 (for last dimension). "
                f"Value was: {channel_dim}"
            )

        self.channel_dim = channel_dim

    def __call__(self, input):
        return F.add_channel(input, channel_dim=self.channel_dim)

    def __repr__(self):
        args = f"(channel_dim={self.channel_dim})"
        return self.__class__.__name__ + args

max_stars_repo_name: jspaaks/vak

max_issues_repo_name: jspaaks/vak

max_forks_repo_name: jspaaks/vak


text:
"""
For the Cov and MAT results of GeoMol
Refer to codes from GeoMol and ConfGF
"""
import numpy as np
import pandas as pd
import argparse
import pickle

from pathlib import Path, PosixPath
from functools import partial
from multiprocessing import Pool
from tqdm import tqdm

from rdkit.Chem.rdForceFieldHelpers import MMFFOptimizeMolecule  # from ConfGF
from rdkit.Chem import rdMolAlign as MA
from rdkit.ML.Descriptors import MoleculeDescriptors
from rdkit.Chem.rdmolops import RemoveHs # We do not remove Hs to show add-hydgeon capability

def evaluate_conf(mols_ref_gen, useFF, threshold, removeH, maxmatches):
    smi, mols_ref, mols_gen = mols_ref_gen
    try:
        if removeH:
            rmsd_confusion_mat = get_rmsd_confusion_matrix_woh(mols_ref, mols_gen, useFF=useFF, maxmatches=maxmatches)
        else:
            rmsd_confusion_mat = get_rmsd_confusion_matrix(mols_ref, mols_gen, useFF=useFF, maxmatches=maxmatches)
        rmsd_ref_min_recall = rmsd_confusion_mat.min(-1)
        rmsd_ref_min_precision = rmsd_confusion_mat.min(0)
        return (
            (rmsd_ref_min_recall <= threshold).mean(),
            rmsd_ref_min_recall.mean(),
            (rmsd_ref_min_precision <= threshold).mean(),
            rmsd_ref_min_precision.mean(),
        )

    except RuntimeError:
        print(smi)
        return 0, 5, 0, 5

def group_conf_rmsd(mols_ref_gen, useFF, removeH, maxmatches):
    smi, _, mols_gen = mols_ref_gen
    group_rmsd = []
    # try:
    if not removeH:
        for i in range(len(mols_gen) - 1):
            for j in range(i + 1, len(mols_gen)):
                group_rmsd.append(MA.GetBestRMS(mols_gen[i], mols_gen[j], maxMatches=maxmatches))
    else:
        for i in range(len(mols_gen) - 1):
            for j in range(i + 1, len(mols_gen)):
                group_rmsd.append(MA.GetBestRMS(RemoveHs(mols_gen[i]), RemoveHs(mols_gen[j]), maxMatches=maxmatches))
    return sum(group_rmsd) / len(group_rmsd)

def get_rmsd_confusion_matrix(mols_ref, mols_gen, useFF, maxmatches):

    num_conf_ref = len(mols_ref)
    num_conf_gen = len(mols_gen)
    rmsd_confusion_mat = -1 * np.ones([num_conf_ref, num_conf_gen], dtype=np.float64)

    for i, conf_ref in enumerate(mols_ref):
        for j, conf_gen in enumerate(mols_gen):
            rmsd_confusion_mat[i, j] = MA.GetBestRMS(
                conf_gen, conf_ref, maxMatches=maxmatches
            )  # maxMatches is important for time
    return rmsd_confusion_mat

def get_rmsd_confusion_matrix_woh(mols_ref, mols_gen, useFF, maxmatches):

    num_conf_ref = len(mols_ref)
    num_conf_gen = len(mols_gen)
    rmsd_confusion_mat = -1 * np.ones([num_conf_ref, num_conf_gen], dtype=np.float64)

    for i, conf_ref in enumerate(mols_ref):
        for j, conf_gen in enumerate(mols_gen):
            rmsd_confusion_mat[i, j] = MA.GetBestRMS(
                RemoveHs(conf_gen), RemoveHs(conf_ref), maxMatches=maxmatches
            )  # maxMatches is important for time
    return rmsd_confusion_mat

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Get GeoMol results")
    parser.add_argument(
        "--ref",
        default="test_ref.pickle",
        type=PosixPath,
        help="absolute path of referenced confs' pickle",
    )
    parser.add_argument(
        "--file", type=str, help="name of tested pickle file",
    )
    parser.add_argument(
        "--testpath", type=PosixPath, help="path of testing, which pkl files exist"
    )
    parser.add_argument(
        "--threshold", type=float, default=0.5, help="rmsd threshould, default to be"
    )
    parser.add_argument(
        "--FF",
        action="store_true",
        default=False,
        help="choose to use FF optimization from RDkit",
    )
    parser.add_argument(
        "--core", type=int, default=2, help="cpu cores used for computation"
    )
    parser.add_argument("--no", type=int, help="analyze the Nth molecule")
    parser.add_argument(
        "--removeH",
        action="store_true",
        default=False,
        help="choose to remove hydrogens for better RMSD",
    )
    parser.add_argument(
        "--rdkerrtxt",
        type=str, help="name of rdkit error smiles .txt file",
        default="rdkit_err_smiles.txt"
    )
    parser.add_argument(
        "--geoerrtxt",
        type=str, help="name of GeoMol error smiles .txt file",
        default="GeoMol_err_smiles.txt"
    )
    parser.add_argument(
        "--maxmatches",
        type=int, help="maxmatches for rdkit align",
        default=100
    )
    parser.add_argument(
        "--datatype",
        type=str, help="['train']",
        default=100
    )

    args = parser.parse_args()

    testfile = args.testpath / args.file
    ref_file = args.testpath / args.ref
    dirty_rdk_smis = args.testpath / args.rdkerrtxt
    dirty_GM_smis = args.testpath / args.geoerrtxt

    with open(ref_file, "rb") as f:
        data_ref = pickle.load(f)
    # dirty_smi_list = dirty_rdk_smis.read_text().split("\n") + dirty_GM_smis.read_text().split("\n") + ["CO[C@H]1[C@H]2CN3C[C@@H]1[C@@H]23"] + ['C[C@@H](O)[C@@H]1[C@H](O)[C@@H]2N[C@@H]21'] + ["'C[C@H]1[C@@H](C=O)N1C'"]
    dirty_smi_list = dirty_rdk_smis.read_text().split("\n") + dirty_GM_smis.read_text().split("\n")
    pklname = testfile.name[:-7]
    print(f"Dealing with {pklname}")

    with open(testfile, "rb") as f:
        data_gen = pickle.load(f)

    covs_R, covs_P = [], []
    mats_R, mats_P = [], []
    group_rmsd = []

    # data_gen_ref = [ (smi, mols_ref, data_gen[smi]) for smi,mols_ref in data_ref.items() if not smi in dirty_smi_list]
    data_gen_ref = [
        (smi, mols_ref, data_gen[smi])
        # (smi, [mols_ref], data_gen[smi]) # For single
        for smi, mols_ref in data_ref.items()
        if not smi in dirty_smi_list
    ]  # for platinum, where each mol has one conf

    data_gen_ref = [data_gen_ref[args.no - 1]] if args.no else data_gen_ref
    func1 = partial(
        evaluate_conf, useFF=args.FF, threshold=args.threshold, removeH=args.removeH, maxmatches=args.maxmatches
    )  # just for imap to take only one positional arguments

    func2 = partial(
        group_conf_rmsd, useFF=args.FF, removeH=args.removeH, maxmatches=args.maxmatches
    )

    with Pool(args.core) as pool:
        for result in tqdm(pool.imap(func1, data_gen_ref), total=len(data_gen_ref)):

            covs_R.append(result[0])
            mats_R.append(result[1])
            covs_P.append(result[2])
            mats_P.append(result[3])

        covs_R, mats_R, covs_P, mats_P = map(np.array, [covs_R, mats_R, covs_P, mats_P])

        for result in tqdm(pool.imap(func2, data_gen_ref), total=len(data_gen_ref)):
            group_rmsd.append(result)
        group_rmsd = np.array(group_rmsd)

    # # pool.close()
    # # pool.join()
    # # pool = multiprocessing.Pool(args.core)

    # for i, mol_gen_ref in tqdm(enumerate(data_gen_ref), total=len(data_gen_ref)):
        # group_rmsd.append(group_conf_rmsd(mol_gen_ref, useFF=args.FF, removeH=args.removeH))
    # group_rmsd = np.array(group_rmsd)

    # Write summary csv
    calculator = MoleculeDescriptors.MolecularDescriptorCalculator(['NumRotatableBonds'])
    numrots = [ calculator.CalcDescriptors(mol[0])[0] for _, mol in data_ref.items() ]
    # numrots = [ calculator.CalcDescriptors(mol)[0] for smi, mol in data_ref.items() if smi not in dirty_smi_list] # For single
    smis = list(data_ref.keys())
    indexes = [smis.index(smi) for smi in smis if smi not in dirty_smi_list]
    num_cov_mat_dict = {
        # 'No.': list(range(1,len(data_ref)+1)),
        'No.': indexes,
        'num_rotatable': numrots,
        'cov_R': covs_R,
        'mat_R': mats_R,
        'cov_P': covs_P,
        'mat_P': mats_P,
    }
    df = pd.DataFrame(num_cov_mat_dict)
################################################################################
    if not args.removeH:
        np.save(args.testpath / f"{pklname}-COV_R-th{args.threshold}-maxm{args.maxmatches}.npy", covs_R)
        np.save(args.testpath / f"{pklname}-MAT_R-th{args.threshold}-maxm{args.maxmatches}.npy", mats_R)
        np.save(args.testpath / f"{pklname}-COV_P-th{args.threshold}-maxm{args.maxmatches}.npy", covs_P)
        np.save(args.testpath / f"{pklname}-MAT_P-th{args.threshold}-maxm{args.maxmatches}.npy", mats_P)
        np.save(args.testpath / f"{pklname}-ingroup-rmsd-maxm{args.maxmatches}.npy", group_rmsd)

        print("Now the threshold is:", args.threshold)
        print("Now max matches for RMSD is:", args.maxmatches)
        print(
            "Coverage Mean (Recall): %.4f | Coverage Median (Recall): %.4f | Match Mean (Recall): %.4f | Match Median (Recall): %.4f"
            % (covs_R.mean(), np.median(covs_R), mats_R.mean(), np.median(mats_R))
        )
        print(
            "Coverage Mean (Precision): %.4f | Coverage Median (Precision): %.4f | Match Mean (Precision): %.4f | Match Median (Precision): %.4f"
            % (covs_P.mean(), np.median(covs_P), mats_P.mean(), np.median(mats_P))
        )

        outfile = args.testpath / f"{pklname}-th{args.threshold}-maxm{args.maxmatches}.csv"
        df.to_csv(outfile, index=False)

    else:
        np.save(args.testpath / f"{pklname}-COV_R-th{args.threshold}-woh-maxm{args.maxmatches}.npy", covs_R)
        np.save(args.testpath / f"{pklname}-MAT_R-th{args.threshold}-woh-maxm{args.maxmatches}.npy", mats_R)
        np.save(args.testpath / f"{pklname}-COV_P-th{args.threshold}-woh-maxm{args.maxmatches}.npy", covs_P)
        np.save(args.testpath / f"{pklname}-MAT_P-th{args.threshold}-woh-maxm{args.maxmatches}.npy", mats_P)
        np.save(args.testpath / f"{pklname}-ingroup-rmsd-woh-maxm{args.maxmatches}.npy", group_rmsd)

        print("Now the threshold is:", args.threshold)
        print("Now max matches for RMSD is:", args.maxmatches)
        print("Note now removeH has been ignited")
        print(
            "Coverage Mean (Recall): %.4f | Coverage Median (Recall): %.4f | Match Mean (Recall): %.4f | Match Median (Recall): %.4f"
            % (covs_R.mean(), np.median(covs_R), mats_R.mean(), np.median(mats_R))
        )
        print(
            "Coverage Mean (Precision): %.4f | Coverage Median (Precision): %.4f | Match Mean (Precision): %.4f | Match Median (Precision): %.4f"
            % (covs_P.mean(), np.median(covs_P), mats_P.mean(), np.median(mats_P))
        )

        outfile = args.testpath / f"{pklname}-th{args.threshold}-removeH-maxm{args.maxmatches}.csv"
        df.to_csv(outfile, index=False)


text:
#!/usr/bin/env python
import yaml
import h5py
import io
import numpy
import os
import pandas
import re
import sklearn.model_selection
import torch
import datasets.padchestmap as padchestmap
from datasets.cxrdataset import CXRDataset, H5Dataset
from PIL import Image

padchesttochexpert = padchestmap.padchesttochexpert
padchesttochexpert['COVID 19'] = 'COVID-19'
padchesttochexpert['COVID 19 uncertain'] = 'COVID-19'
padchesttochexpert['viral pneumonia'] = 'Pneumonia'

padchesttochestxray14 = padchestmap.padchesttochestxray14
padchesttochestxray14['COVID 19'] = 'COVID-19'
padchesttochestxray14['COVID 19 uncertain'] = 'COVID-19'
padchesttochestxray14['viral pneumonia'] = 'Pneumonia'

SKIP_WINDOWING = [
        "sub-S03169_ses-E07193_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03169_ses-E07878_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03067_ses-E07350_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03067_ses-E07908_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03240_ses-E07784_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03218_ses-E07970_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03240_ses-E07626_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03214_ses-E07979_run-1_bp-chest_vp-ap_dx.png"
        ]

FLIP = ["sub-S03936_ses-E08636_run-1_bp-chest_vp-ap_dx.png",
        "sub-S04275_ses-E08736_run-1_bp-chest_vp-ap_dx.png",
        "sub-S04190_ses-E08730_run-1_bp-chest_vp-ap_dx.png",
        "sub-S03068_ses-E06592_run-1_bp-chest_vp-ap_cr.png"]

def grouped_split(dataframe, random_state=None, test_size=0.05):
    '''
    Split a dataframe such that patients are disjoint in the resulting folds.
    The dataframe must have an index that contains strings that may be processed
    by _get_patient_id to return the unique patient identifiers.
    '''
    groups = _get_unique_patient_ids(dataframe)
    traingroups, testgroups = sklearn.model_selection.train_test_split(
            groups,
            random_state=random_state,
            test_size=test_size)
    traingroups = set(traingroups)
    testgroups = set(testgroups)

    traindict = {}
    testdict = {}
    for idx, row in dataframe.iterrows():
        patient_id = _get_patient_id(dataframe, idx)
        if patient_id in traingroups:
            traindict[idx] = row
        elif patient_id in testgroups:
            testdict[idx] = row
    traindf = pandas.DataFrame.from_dict(
        traindict,
        orient='index',
        columns=dataframe.columns)
    testdf = pandas.DataFrame.from_dict(
        testdict,
        orient='index',
        columns=dataframe.columns)
    return traindf, testdf

def _get_patient_id(df, idx):
    return df['participant'].loc[idx]

def _get_unique_patient_ids(dataframe):
    ids = list(dataframe.index)
    ids = [_get_patient_id(dataframe, i) for i in ids]
    ids = list(set(ids))
    ids.sort()
    return ids

class BIMCVCOVIDDataset(H5Dataset, CXRDataset):
    '''
    projection labels include AP, PA, LAT, and AP_SUPINE (based on DICOM data).
      if 'include_unknown_projections == True' then the additional hand-labeled
      projections 'frontal' and 'lateral' will be included.
    '''
    def __init__(self, fold, random_state=30493, include_lateral=False,
                 include_unknown_projections=False, include_ap_supine=False,
                 include_unknown_labels=False, initialize_h5=False, covid_labels='molecular',
                 labels='chexpert', ):
        '''
        covid_labels (str): if 'molecular', the COVID label associated with each
          image will be based on molecular assay results (PCR or serology); all
          images will be COVID positive. If 'radiologic', the COVID label will
          be based on presence of radiological evidence of COVID.
        '''
        self.labelstyle = labels.lower()
        if self.labelstyle == 'chexpert':
            self.labels = [
                'Enlarged Cardiomediastinum',
                'Cardiomegaly',
                'Lung Opacity',
                'Lung Lesion',
                'Edema',
                'Consolidation',
                'Pneumonia',
                'Atelectasis',
                'Pneumothorax',
                'Pleural Effusion',
                'Pleural Other',
                'Fracture',
                'Support Devices',
                'COVID-19']
        elif self.labelstyle == 'chestx-ray14':
            self.labels = [
                'Atelectasis',
                'Cardiomegaly',
                'Consolidation',
                'Edema',
                'Effusion',
                'Emphysema',
                'Fibrosis',
                'Hernia',
                'Infiltration',
                'Mass',
                'Nodule',
                'Pleural_Thickening',
                'Pneumonia',
                'Pneumothorax',
                'COVID-19']
        else:
            raise NotImplementedError
        self.covid_labels = covid_labels.lower()
        if not self.covid_labels in ['radiologic', 'molecular']:
            raise ValueError('Invalid value {:s} for keyword argument "covid_labels"'.format(repr(self.covid_labels)))
        if self.covid_labels == 'radiologic':
            padchesttochestxray14['viral pneumonia'] = 'COVID-19'
            padchesttochexpert['viral pneumonia'] = 'COVID-19'
        self.report_regex = re.compile('ses-E\d+')
        self.datapath = 'data/BIMCV-COVID-19'
        self.labelpath = 'derivatives/labels/labels_covid19_posi.tsv'
        self.unknown_label_path = 'datasets/bimcv_covid_unknown_labels.txt'
        self.labeldf = pandas.read_csv(os.path.join(self.datapath, self.labelpath), delimiter='\t')
        self.h5path = 'data/BIMCV-COVID-19/BIMCV-COVID-19.h5'
        self.df = pandas.read_csv(os.path.join(self.datapath, 'BIMCV-COVID-19.csv'))
        # Filter images with windowing data
        self.df = self.df.query('window_center == window_center | lut == lut') # remove NaN
        self.df.lut = self.df.lut.apply(lambda x: eval(x) if isinstance(x,str) else x) # strings to LUT lists
        if include_unknown_projections: # unlabeled projections
            with open("datasets/bimcv_covid_manual_projection_labels.yml", 'r') as handle:
                projection_labels = yaml.load(handle)
            unknowns = self.df.query("projection == 'UNK'")
            for idx, row in unknowns.iterrows():
                self.df.projection.loc[idx] = projection_labels[row.path]
        else:
            self.df = self.df.query("projection == 'AP' or projection == 'PA' or projection == 'LAT' or projection == 'AP SUPINE'")
        if not include_lateral: # lateral projections
            self.df = self.df.query("projection == 'AP' or projection == 'PA' or projection == 'frontal' or projection == 'AP SUPINE'")
        if not include_ap_supine:
            self.df = self.df.query("projection == 'AP' or projection == 'PA' or projection == 'frontal' or projection == 'LAT'")
        if not include_unknown_labels:
            with open(self.unknown_label_path, 'r') as handle:
                unknown_paths = []
                for line in handle:
                    unknown_paths.append(line.strip())
                self.df = self.df[~self.df['path'].isin(unknown_paths)]
        if not fold in ['train', 'val', 'test', 'all']:
            raise ValueError("Invalid fold: {:s}".format(fold))
        self._transforms['all'] = self._transforms['val']
        self.transform = self._transforms[fold]
        if fold == 'all':
            pass
        else:
            trainvaldf, testdf = grouped_split(
                    self.df,
                    random_state=random_state,
                    test_size=0.05)
            if fold == 'train' or fold == 'val':
                traindf, valdf = grouped_split(
                        trainvaldf,
                        random_state=random_state,
                        test_size=0.05)
                if fold == 'train':
                    self.df = traindf
                elif fold == 'val':
                    self.df = valdf
            elif fold == 'test':
                self.df = testdf
        if initialize_h5:
            self.init_worker(None)

    def init_worker(self, worker_id):
        self.h5 = h5py.File(self.h5path, 'r', swmr=True)

    def __getitem__(self, idx):
        UINT8_MAX = 255

        image = self._raw_image_from_disk(idx)
        imagename = self.df.path.iloc[idx].split('/')[-1]
        if imagename in SKIP_WINDOWING:
            pass
        else:
            image = numpy.array(image, dtype=numpy.int64)
            # Use LUT if we have it
            lut = self.df.lut.iloc[idx]
            if isinstance(lut, list):
                lut_min = int(self.df.lut_min.iloc[idx])
                lut = numpy.array(lut)
                lut = numpy.concatenate((numpy.ones(lut_min)*lut[0],
                                         lut,
                                         numpy.ones(65536-lut_min-len(lut))*lut[-1]), axis=0)
                image = lut[image]
                if self.df.rescale_slope.iloc[idx]:
                    image *= self.df.rescale_slope.iloc[idx] + self.df.rescale_intercept.iloc[idx]
                max_ = 2**self.df.bits_stored.iloc[idx]-1
                image = image.astype(numpy.float64)/max_
            else: # use window data
                window_center = self.df.window_center.iloc[idx]
                window_width = self.df.window_width.iloc[idx]
                window_min = int(window_center - window_width/2)
                image -= window_min
                image = image.astype(numpy.float64)*1/window_width
            # clip
            image[image<0] = 0
            image[image>1] = 1
            image = (image*UINT8_MAX).astype(numpy.uint8)

            # invert if needed
            photometric_interpretation = self.df.photometric_interpretation.iloc[idx]
            if photometric_interpretation == 'MONOCHROME1':
                image = UINT8_MAX-image
            elif photometric_interpretation == 'MONOCHROME2':
                image = image
            else:
                raise ValueError('unknown photometric interpretation: {:s}'.format(photometric_interpretation))
            if imagename in FLIP:
                image = numpy.flipud(image)
            try:
                image = Image.fromarray(image, mode='L')
            except ValueError:
                image = Image.fromarray(image, mode='RGB')
        image = image.convert('RGB')

        if self.transform:
            image = self.transform(image)
        label = self.get_labels(idx)
        return image, label, 0, 0

    def _raw_image_from_disk(self, idx):
        '''
        Retrieve the raw PIL image from storage.
        '''
        imagename = self.df.path.iloc[idx].split('/',1)[1]
        data = self.h5['images'].get(imagename)
        image = Image.open(io.BytesIO(numpy.array(data)))
        return image

    def get_labels(self, idx):
        '''
        Get the labels for index ``idx``.

        An initial draft of the labels is obtained via _parse_labels, which is
        then converted to the desired style of labels.
        '''
        findings = self._parse_labels(idx)
        if self.covid_labels == 'molecular':
            # ALL samples are covid-19+
            if not 'COVID 19' in findings:
                findings.append('COVID 19')

        if self.labelstyle == 'chexpert':
            map_ = padchesttochexpert
        elif self.labelstyle == 'chestx-ray14':
            map_ = padchesttochestxray14
        else:
            raise ValueError
        labels = []
        for f in findings:
            # map the PadChest finding to the CheXpert/chestx-ray14 label
            if f == '':
               continue
            converted = map_[f]
            if converted == '':
               continue
            elif isinstance(converted, str):
                labels.append(converted)
            elif isinstance(converted, list):
                labels += converted
            else:
                raise ValueError(str(converted))
        labels = set(labels)

        # convert to one-hot encoding
        label = numpy.zeros(len(self.labels), dtype=int)
        for i in range(0, len(self.labels)):
            if self.labels[i] != "N/A" and self.labels[i] in labels:
                label[i] = 1
        return label

    def _parse_labels(self, idx):
        '''
        This requires a few steps:
        1) The Labels field of the tsv file contains a string, which can
           possibly be evaluated to a Python list, or `NaN`.
        2) Handle the label `unchanged`. Look to the most recent previous study
           and copy over the labels. This may have to be done recursively, and
           sometimes the oldest study may not contain labels. In this case, we
           give the study no labels.
        '''
        path = self.df.path.iloc[idx]
        reportid = self.report_regex.search(path).group()
        # s will be something like
        #'["costophrenic angle blunting", "loculated pleural effusion", "clavicle fracture"]'
        row = self.labeldf.query("ReportID == @reportid")
        try:
            labeldf_idx = row.index[0]
        except IndexError: # no report found!
            #with open('temp.txt', 'a') as handle:
            #    handle.write(path + '\n')
            return []
        s = row.Labels.iloc[0]
        # convert the string to a list
        if isinstance(s, str):
            findings = [substr.strip().strip("'").strip() for substr in s.strip("[]").split('\t')]
        else:
            findings = []

        if 'unchanged' in findings:
            # Look backward toward the most recent study
            patientid = self.labeldf.PatientID[labeldf_idx]
            study_date = self.df.study_date.iloc[idx] # numpy.int64
            study_time = self.df.study_time.iloc[idx]
            other_studies = self.df.query(''' participant == "{:s}" '''.format(patientid))
            previous_studies = other_studies.query("study_date < @study_date | (study_date == @study_date & study_time < @study_time)")
            if len(previous_studies) > 0:
                other_studies = previous_studies.sort_values(by="study_time", kind='mergesort')
                # stable sort here to previous time ordering
                most_recent_study = previous_studies.sort_values(by="study_date", kind='mergesort').iloc[-1]
                findings.remove("unchanged")
                new_idx = self.df.index.get_loc(most_recent_study.name)
                findings += self._parse_labels(new_idx)
            else:
                findings.remove("unchanged")
        return findings

    def get_all_labels(self):
        arr = numpy.zeros((len(self), len(self.labels)))
        for i in range(len(self)):
            arr[i] = self.get_labels(i)
        return arr


text:
import numpy as np

from argoverse.utils.metric_time import TimeUnit, to_metric_time

def test_nanoseconds_to_seconds():
    """Test conversion from nanoseconds to seconds."""
    ts_ns = 950000000
    ts_s = to_metric_time(ts=ts_ns, src=TimeUnit.Nanosecond, dst=TimeUnit.Second)
    assert np.isclose(ts_s, 0.95)
    assert np.isclose(ts_s, ts_ns / 1e9)

def test_seconds_to_nanoseconds():
    """Test conversion from seconds to nanoseconds."""
    ts_s = 0.95
    ts_ns = to_metric_time(ts=ts_s, src=TimeUnit.Second, dst=TimeUnit.Nanosecond)
    assert np.isclose(ts_ns, 950000000)
    assert np.isclose(ts_ns, ts_s * 1e9)

def test_milliseconds_to_seconds():
    """Test conversion from milliseconds to seconds."""
    ts_ms = 300
    ts_s = to_metric_time(ts=ts_ms, src=TimeUnit.Millisecond, dst=TimeUnit.Second)
    assert np.isclose(ts_s, 0.3)

def test_seconds_to_milliseconds():
    """Test conversion from seconds to milliseconds."""
    ts_s = 0.3
    ts_ms = to_metric_time(ts=ts_s, src=TimeUnit.Second, dst=TimeUnit.Millisecond)
    assert np.isclose(ts_ms, 300)

max_stars_repo_name: gargrohin/argoverse-api

max_issues_repo_name: gargrohin/argoverse-api

max_forks_repo_name: gargrohin/argoverse-api


text:
import numpy as np
import cv2

def Sobel(img, thres):
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    im_x = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3)
    im_x = cv2.convertScaleAbs(im_x)
    im_y = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3)
    im_y = cv2.convertScaleAbs(im_y)
    img = cv2.addWeighted(im_x, 1, im_y, 1, 0)

    ret, thresh = cv2.threshold(img, thres, 255, cv2.THRESH_BINARY)
    return thresh

def Hough(im1, im2, thres):
    # im2 should be grayscale image
    lines = cv2.HoughLinesP(im2, 1, np.pi/180,thres)
    line_img = np.zeros_like(im1)
    for line in lines:
        for x1,y1,x2,y2 in line:
            cv2.line(line_img,(x1,y1),(x2,y2),[255,0,255],2)

    # blending factor can be adjust
    res = cv2.addWeighted(im1, 0.4, line_img, 0.6, 0)
    return res

def SBM(imL, imR, numdisp, blocksize):
    # Both input images are grayscale
    sbm = cv2.StereoBM_create(numDisparities = numdisp, blocksize = blocksize)
    disparity = sbm.compute(imL, imR)
    return disparity

def feature_flow(im1, im2, feature_num):
    feature_params = dict(maxCorners = feature_num, qualityLevel = 0.01, minDistance = 7, blockSize = 7 )
    lk_params = dict( winSize = (15, 15), maxLevel = 2, criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
    color = np.zeros([feature_num,3])
    color[:] = [200, 0, 200]

    im1_g = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
    im2_g = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY)

    im1_g = cv2.bilateralFilter(im1_g, 3, 50, 50)
    im2_g = cv2.bilateralFilter(im2_g, 3, 50, 50)

    p0 = cv2.goodFeaturesToTrack(im1_g, mask=None, **feature_params)
    mask = np.zeros_like(im1)
    p1, st, err = cv2.calcOpticalFlowPyrLK(im1_g, im2_g, p0, None,  **lk_params)

    good_new = p1[st == 1]
    good_old = p0[st == 1]

    for i, (new, old) in enumerate(zip(good_new, good_old)):
        a, b = new.ravel()
        c, d = old.ravel()
        mask = cv2.line(mask, (a, b), (c, d), color[i].tolist(), 2)
        right = cv2.circle(im2, (a, b), 5,  color[i].tolist(), -1)

    img = cv2.add(im2, mask)
    return img

def grid_flow(im_g, im, flow, black = False):
    # im : grayscale image
    step = 20
    h, w = im_g.shape[:2]
    y, x = np.mgrid[step/2:h:step, step/2:w:step].reshape(2,-1).astype(int)
    fx, fy = flow[y,x].T

    lines = np.vstack([x, y, x+fx, y+fy]).T.reshape(-1,2,2)
    lines = np.int32(lines + 0.5)
    '''vis = cv2.cvtColor(im_g, cv2.COLOR_GRAY2BGR)
    if black:
        vis = np.zeros((h, w, 3), np.uint8)'''
    cv2.polylines(im, lines, 0, (0, 255, 0))

    for (x1, y1), (x2, y2) in lines:
        cv2.circle(im, (x1, y1), 1, (200, 0, 200), -1)
    return im

def Min():
    print("Slave 511... so sad")