import argparse
import os,glob,sys,yaml,cv2
import time,datetime
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.nn.functional as F
import torch.utils.data
from thop import profile
INVALID_DISPARITY_F32 = 0.-300.
INVALID_DISPARITY_U16 = 32768

from util import save_checkpoint,AverageMeter,save_float_to_colormap,normerWrong,normer,normerImageNet,resizer,constantNorm
# from  lightStereo.lightstereo_refine  import LightStereo 
from  lightStereo.lightstereo  import LightStereo 

def dispFloatToUint16(dispFOrig, ratio=1.):
    '''
    disparity range from -300 to -0 is currently not suppoted for read/write to png file
    '''
    dispF     = dispFOrig.copy()*ratio
    dispInvMsk=(dispF<=INVALID_DISPARITY_F32*ratio)
    dispValMsk=(dispF>INVALID_DISPARITY_F32*ratio)
    dispF[dispValMsk] *= 128.
    dispF[dispInvMsk]  = INVALID_DISPARITY_U16
    dispU16 = dispF.astype(np.uint16)
    return dispU16

def convert_relu_to_relu6(model):
    for child_name, child in model.named_children():
        if isinstance(child, (nn.ReLU, nn.LeakyReLU)):
            #print(child_name)
            setattr(model, child_name, nn.ReLU6(inplace=True))
        else:
            convert_relu_to_relu6(child)

# gaussian blur
def gaussBlurCore(img, ksize=(3,3), sigma=0.35):
    kw = int(ksize[0])
    kh = int(ksize[1])
    blurImg = cv2.GaussianBlur(img, (kw,kh), sigma)
    return blurImg

def gaussBlur(left, right):
    left, right = left.copy(), right.copy()
    left, right = left, right

    left     = gaussBlurCore(left)
    left,right   = left, right
    return left, right

def colorTransfer(Isrc, Iref):
    Is = Isrc.copy()
    Ir = Iref.copy()
    #plt.imshow(Ir[...,::-1])   
    #plt.show()
    #plt.imshow(Is[...,::-1])
    #plt.show()
    #BGR->LAB
    valid = ((Iref[:,:,0]>0) & (Iref[:,:,1]>0) & (Iref[:,:,2]>0))
    LabIs = cv2.cvtColor(Is,cv2.COLOR_BGR2LAB)
    LabIr = cv2.cvtColor(Ir,cv2.COLOR_BGR2LAB)
    #mean、std
    Is_means = [0, 0, 0]
    Ir_means = [0, 0, 0]
    Is_stdevs = [0, 0, 0]
    Ir_stdevs = [0, 0, 0]
    LabIs = LabIs.astype(np.float32)/255.
    LabIr = LabIr.astype(np.float32)/255.
    for i in range(3):
        Is_means[i] += LabIs[:,:,i].mean()
        Ir_means[i] += LabIr[:,:,i][valid].mean()
        Is_stdevs[i] += (LabIs[:,:,i].std() + 1e-9)
        Ir_stdevs[i] += (LabIr[:,:,i][valid].std() + 1e-9)
    thresh = [a/b for a,b in zip(Ir_stdevs,Is_stdevs)]
    LabIt = thresh*(LabIs -Is_means) + Ir_means

    # [0-255]
    LabIt = (LabIt*255.)
    LabIt *= (LabIt>0)
    LabIt = (LabIt * (LabIt<=255) + 255 * (LabIt>255)).astype(np.uint8)

    #show
    It = cv2.cvtColor(LabIt,cv2.COLOR_LAB2BGR)
    #plt.imshow(It)
    #plt.show()

    return It

def getYamlCfg(fn='models/configModelTest.yaml'):
    f       = open(fn, 'r', encoding='utf-8')
    cfg     = f.read()
    content = yaml.load(cfg, Loader=yaml.FullLoader)
    return content

def initModels(content, modelFolder='models', args=None):
    numOfModels = len(list(content.keys()))
    #importedFolder= __import__(modelFolder)

    singleModel  = {}
    for idx in range(numOfModels):
        modelName    = list(content.keys())[idx]
        if content[modelName]['used4Train'] is not True:
            continue

        print('=> Importing model-{}'.format(modelName))
        importedFolder = __import__(content[modelName]['modelFolder'])
        moduleImported = getattr(importedFolder, content[modelName]['moduleName'])
        classImported  = getattr(moduleImported, content[modelName]['className'])

        singleModel['classInst']   = classImported(batchNorm=content[modelName]['batchNorm'], pred_disp=content[modelName]['pred_disp'], args=args)
        break

    return singleModel['classInst']

def initModelsBug(content, modelFolder='models'):
    numOfModels = len(list(content.keys()))
    importedFolder= __import__(modelFolder)

    singleModel  = {}
    for idx in range(numOfModels):
        modelName    = list(content.keys())[idx]
        if content[modelName]['used4Train'] is not True:
            continue

        print('=> Importing model-{}'.format(modelName))
        moduleImported = getattr(importedFolder, content[modelName]['moduleName'])
        classImported  = getattr(moduleImported, content[modelName]['className'])

        singleModel['classInst']   = classImported(batchNorm=content[modelName]['batchNorm'], pred_disp=content[modelName]['pred_disp'])
        break

    return singleModel['classInst']

def loadPretrain(model, pretrained, strict=False):
    oldCkpt = torch.load(pretrained)
    if pretrained.find('.tar')>=0:
        oldCkpt = oldCkpt['state_dict']
    newCkpt = {}
    for k in oldCkpt.keys():
        newCkpt[k.replace('module.','')] = oldCkpt[k]
    model.load_state_dict(newCkpt, strict=strict)
    return model

# Training settings
parser = argparse.ArgumentParser(description='Stereo training...')
parser.add_argument('--data_parent_path', type=str, default='/home/notebook/data/group/xizhong_xxz/', help="location to save models")
parser.add_argument('--data_config_path', type=str, default='./configs/configDatasetMultiscale.yaml', help="location to save models")
parser.add_argument('--model_config_path', type=str, default='./configs/configModel.yaml', help="location to save models")
parser.add_argument('--loss_config_path', type=str, default='./configs/configLoss.yaml', help="location to save models")
parser.add_argument('--save_path', type=str, default='./rectify', help="location to save models")
parser.add_argument('--use_gru_refine', type=int, default=0, help="if use gru refine for disp output")
parser.add_argument('--useBackboneV2', type=int, default=0, help='1 use mobilenet')
parser.add_argument('--att_relu6', type=int, default=0, help='1 use relu6 for att module')
parser.add_argument('--convex_up', type=int, default=0, help='1 use convex upsample')
parser.add_argument('--transAtt', type=int, default=0, help='1 use self att')
parser.add_argument('--agg_blocks', default=[8,16,32], metavar='N', nargs='*', help='set the block number for each stage')
parser.add_argument('--pretrained_Backbone', type=int, default=1, help='1 use pretrained backbone to extractor feature')
parser.add_argument('--sample_tail', type=int, default=0,  help="if add a tail refinement for output quantize")
parser.add_argument('--use_mono', type=int, default=0,  help="if use mono optimize stereo")

parser.add_argument('--useResNet', type=int, default=0, help='use resnet as backbone')
parser.add_argument('--validateOnly', type=int, default=0, help='do validation only. Default=0')
parser.add_argument('--warmup', type=int, default=0, help='do warmup training or not, Default=0')
parser.add_argument('--abnormAug', type=int, default=0, help='do abnormal augmenation such as thin plate transform or diagonal epipolar transform')
parser.add_argument('--eraseAug', type=int, default=0, help='do erasing augmenation')
parser.add_argument('--occAug', type=int, default=0, help='do occlusion augmentation in right image branch')
parser.add_argument('--defocusAug', type=int, default=0, help='do blur augmentation in all image branches')
parser.add_argument('--gradLoss', type=int, default=0, help='use gradient loss for supervision')
parser.add_argument('--normLoss', type=int, default=0, help='use normal loss for supervision')
parser.add_argument('--smtLoss', type=int, default=0, help='use smoothness loss for self-supervision')
parser.add_argument('--use_dn', type=int, default=0, help='0 for bn, 1 for dn')
parser.add_argument('--mobile', type=int, default=0, help='0 for bn, 1 for mobile')
parser.add_argument('--use_relu6', type=int, default=0, help='0 for relu, 1 for relu6')
parser.add_argument('--use_refine_relu6', type=int, default=0, help='0 for relu, 1 for relu6')
parser.add_argument('--backbone_relu6_only', type=int, default=0, help='0 for relu, 1 for relu6')
parser.add_argument('--noNormInput', type=int, default=0, help='normalize input to 0.0 ~ 1.0')
parser.add_argument('--rgbOrder', type=int, default=0, help='0 for bgr, 1 for rgb')
parser.add_argument('--normGT', type=int, default=0, help='normalize GT to 0.0 ~ 1.0')
parser.add_argument('--ada_sigma', type=int, default=0, help='0 for fixed sigma, 1 for adaptive sigma')
parser.add_argument('--depth_wise', type=int, default=0, help='depth wise conv for decoder')
parser.add_argument('--ada_cost', type=int, default=1, help='0 for fixed cost, 1 for adaptive cost')
parser.add_argument('--simple_refine', type=int, default=0, help='1 for simple refiner, 0 for warp refiner')
parser.add_argument('--rightDown2x', type=int, default=0, help='1 for right image plane down 2x, 0 for no downscale')
parser.add_argument('--clipGrad', type=int, default=0, help='use gradient clipping for training')
parser.add_argument('--clip', type=float, default=2.1, help='clipping range')
parser.add_argument('--skyFill', type=int, default=1, help='fill sky area with disparity 0.0 with semantic label')
parser.add_argument('--use_refine', type=int, default=0, help='use refinement blocks')
parser.add_argument('--refine_stages', type=int, default=2, help='use refinement blocks')
parser.add_argument('--warp_feature', type=int, default=0, help='warp feature or original image')
parser.add_argument('--wRec', type=float, default=0.0, help='gradient loss weight')
parser.add_argument('--proj_ch', type=int, default=3, help='channels to reconsrtuct')
parser.add_argument('--pre_scale', type=float, default=0.666667, help='pre-scale')
parser.add_argument('--cat2Fuse', type=int, default=0, help='costs cat to fuse')
parser.add_argument('--trainSceneFlowOnly', type=int, default=0, help='train scene flow dataset only')
parser.add_argument('--wGrad', type=float, default=0.5, help='gradient loss weight')
parser.add_argument('--wSmt', type=float, default=0.5, help='smooth loss weight')
parser.add_argument('--inconsistAug', type=int, default=0, help='inconsistancy augmentation for LR pairs')
parser.add_argument('--imageNetInputNorm', type=int, default=0, help='inputs normalized like in imagenet')
parser.add_argument('--constantNorm', type=int, default=0, help='inputs normalized to [-1, 1]')
parser.add_argument('--useTransformer', type=int, default=0, help='use transformer blocks')
parser.add_argument('--correctCorr', type=int, default=0, help='use correct correlation')
parser.add_argument('--copyPaste', type=int, default=1, help='use copy paste tech.')
parser.add_argument('--zeroBlack', type=int, default=0, help='black area filled with zero when normalizing.')
parser.add_argument('--zeroInvalidGt', type=int, default=0, help='black area filled with zero gt.')

parser.add_argument('--maxdisp', type=int, default=120, help="max disparity range")
#parser.add_argument('--maxdisp', type=int, default=120, help="max disparity range")

parser.add_argument('--shift', type=int, default=0, help='random shift of left image. Default=0')
#shift value must be zero for cost volume aggregation (soft argmin) based network
parser.add_argument('--crop_height', type=int, default=640, help="crop height")
parser.add_argument('--crop_width', type=int, default=768,  help="crop width")
parser.add_argument('--resume', type=str, default='', help="resume from saved model")
parser.add_argument('--start_epoch', type=int, default=0, help='the epoch to resume training')
parser.add_argument('--pretrained', dest='pretrained', default=None, help='path to pre-trained model')

parser.add_argument('--batch_size', type=int, default=8, help='training batch size')
parser.add_argument('--val_batch_size', type=int, default=8, help='testing batch size')
parser.add_argument('--workers', type=int, default=32, help='number of threads')
parser.add_argument('--nEpochs', type=int, default=64, help='number of epochs to train for')
parser.add_argument('--solver', default='adam',choices=['adam','sgd'], help='solver algorithms')
parser.add_argument('--lr', type=float, default=0.0008, help='Learning Rate. Default=0.001')
parser.add_argument('--lr_decay', type=float, default=0.5, help='Learning Rate descending rete. Default=0.5')

'''
Attention: milestones should be ignored(keep unchanged), since it's adapted to nEpochs in the latter main function
'''
#####parser.add_argument('--milestones', default=[14,28,42,56], metavar='N', nargs='*', help='epochs at which learning rate is divided by 2')   
parser.add_argument('--milestones', default=[10,20,30,40,50,60], metavar='N', nargs='*', help='epochs at which learning rate is divided by 2')
#parser.add_argument('--milestones', default=[13,23,33,43,53,63], metavar='N', nargs='*', help='epochs at which learning rate is divided by 2')

#parser.add_argument('--weight_decay', '--wd', default=4e-4, type=float, metavar='W', help='weight decay')
parser.add_argument('--weight_decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay')
parser.add_argument('--bias_decay', default=0, type=float, metavar='B', help='bias decay')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum for sgd, alpha parameter for adam')
parser.add_argument('--beta', default=0.999, type=float, metavar='M', help='beta parameter for adam')


'''
ascending order means use loss.sum()/batchSize, which is map level average
descending order means use loss.mean(), which is pixel level average
'''
#parser.add_argument('--multiscale_weights', '-w', default=[0.005,0.01,0.02,0.08,0.32], type=float, nargs=5, help='training weight for each scale, from highest resolution (flow2) to lowest (flow6)', metavar=('W2', 'W3', 'W4', 'W5', 'W6'))
#parser.add_argument('--multiscale_weights', '-w', default=[0.32,0.08,0.02,0.01,0.005], type=float, nargs=5, help='training weight for each scale, from highest resolution (flow2) to lowest (flow6)', metavar=('W2', 'W3', 'W4', 'W5', 'W6'))
parser.add_argument('--multiscale_weights', '-w', default=[1./3, 2./3, 1., 1., 1.], type=float, nargs=5, help='training weight for each scale, from highest resolution (flow2) to lowest (flow6)', metavar=('W2', 'W3', 'W4', 'W5', 'W6'))

parser.add_argument('--seed', type=int, default=2022, help='random seed to use. Default=123')

parser.add_argument('--print_freq', '-p', default=50, type=int, metavar='N', help='print frequency')
parser.add_argument('--save_freq', '-s', default=5, type=int, metavar='N', help='save checkpoint frequency')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def main():

    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    print('=> get test data pairs')
    # args.data_path="/home/notebook/data/group/liulibo/stereoProject/rectify/bars/"
    # args.data_path="/home/notebook/data/group/liulibo/MiDaS/langan_refine_data/"
    args.data_path="/home/notebook/data/group/Ruby/langan_datasets/shoushi/shoushi_crop_s/mono/"

    args.flip_LR=0
    args.useMinMax=0
    lists  =[args.data_path+file  for file in os.listdir(args.data_path)  if(("left"  in file  or "master"  in file)  and "disp"  not in file)]
    total  = len(lists)
    print(total)
    #pause

    print('=> get models from configuration file')
    content = getYamlCfg(fn='configs/configModelTest.yaml')
    model = LightStereo(parser)
    model = model.to(device).eval()
    # model   = loadPretrain(model, "./checkpoint/refiner_finetune_20240918/model_best.pth.tar", strict=False)
    model   = loadPretrain(model, "./checkpoint/gru/best.pth.tar", strict=False)

    print('=> infer')
    count=0
    # for idx in range(total):
    for file in os.listdir(args.data_path):  
        if(("left"  in file  or "master"  in file)):
            if "Disp"   in file:
                continue
            try:
                srcLFn    = args.data_path+file
                srcRFn    = srcLFn.replace('left','right').replace('master', 'slave')
                print(srcLFn)
                print(srcRFn)
                lImg      = cv2.imread(srcLFn,1)
                rImg      = cv2.imread(srcRFn,1)
                h,w,c=lImg.shape
                hh,ww,cc=rImg.shape
                lMask     = ((lImg[:,:,0]==0) & (lImg[:,:,1]==0) & (lImg[:,:,2]==0)).astype(np.uint8)
            except:
                continue
            if args.flip_LR>0:
                lImg  = lImg[:,::-1,:]
                rImg  = rImg[:,::-1,:]
                lMask = lMask[:,::-1]

            rszH,rszW = lImg.shape[0],lImg.shape[1]
            lImg,rImg = resizer(lImg,rImg,args)

            # if args.rgbOrder > 0:
            #     lImg= lImg[:,:,::-1].copy()
            #     rImg= rImg[:,:,::-1].copy()
            # if args.imageNetInputNorm <= 0:
            #     imgPair   = normer(lImg, rImg)
            # else:
            #     if args.constantNorm:
            #         imgPair   = constantNorm(lImg, rImg)
            #     else:
            imgPair   = normerImageNet(lImg, rImg)
            
            imgPairTs = torch.from_numpy(imgPair).to(device).float().unsqueeze(0)

            with torch.no_grad():        
                output= model(imgPairTs)

            if isinstance(output, list):
                output    = output[-1].squeeze().cpu().numpy()
            else:
                output    = output.squeeze().cpu().numpy()
            if args.flip_LR>0:
                output    = output[:,::-1]
                lMask     = lMask[:,::-1]

            outH,outW = output.shape[0],output.shape[1]
            scaleDisp = rszW * 1. / outW

            output    = cv2.resize(output, (rszW,rszH), interpolation=cv2.INTER_LINEAR) * scaleDisp
            #print(output.min())

            # dispU16   = dispFloatToUint16(output, ratio=3)
            output_path = '/home/notebook/data/group/Ruby/langan_datasets/shoushi/shoushi_crop_s/mono/'
            cv2.imwrite(output_path +'/'+os.path.basename(srcLFn).replace('left','pred').replace('master','pred').replace('.jpg','.png'), output*10)
            # cv2.imwrite(args.data_path+'/'+os.path.basename(srcLFn).replace('left','pred').replace('master','pred').replace('.jpg','.png'), output*10)

            # print(args.save_path)
            # cv2.imwrite(args.data_path+'/'+os.path.basename(srcLFn).replace('left','gru_leftDisp_gray').replace('master','gru_leftDisp_gray').replace('.jpg','.png'), dispU16)
            # color     = save_float_to_colormap(output, saveFn=args.data_path+'/'+os.path.basename(srcLFn).replace('left','gru_leftDisp').replace('master','gru_leftDisp').replace('.jpg','.png'), inputLimMin=0, inputLimMax=120, useMinMaxFirst=(args.useMinMax>0), saveU8=False)
            # cv2.imwrite(args.data_path+'/'+os.path.basename(srcLFn).replace('left','leftDisp_gray').replace('master','leftDisp_gray').replace('.jpg','.png'), dispU16)
            # color     = save_float_to_colormap(output, saveFn=args.data_path+'/'+os.path.basename(srcLFn).replace('left','leftDisp').replace('master','leftDisp').replace('.jpg','.png'), inputLimMin=0, inputLimMax=120, useMinMaxFirst=(args.useMinMax>0), saveU8=False)
        
if __name__ == '__main__':
    main()

对于这个代码，写一个能够不处理已经处理过的文件的功能