import argparse
import time
import logging
import tqdm
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.utils.data
from torch.utils.data import Dataset
import cv2
from torch.utils.data import DataLoader
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import numpy as np
import PIL
from PIL import Image
from PIL import ImageOps
import os
from repvit  import Mono_Depth
from efficientnet_lite_concat_v1 import efficientnet_lite_params,build_efficientnet_lite

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class AuxModule(nn.Module):
    def __init__(self, model):
        super(AuxModule, self).__init__()
        self.module = model

    def forward(self, rgb):
        return self.module(rgb)
def randomPerspectiveTransformCore(image, maxDimRatio=1./20):
    rows, cols, _ = image.shape
    
    dxLt = np.random.randint(0, cols * maxDimRatio)
    dyLt = np.random.randint(0, rows * maxDimRatio)
    dxRt = np.random.randint(-cols * maxDimRatio, 0)
    dyRt = np.random.randint(0, rows * maxDimRatio)
    dxBr = np.random.randint(-cols * maxDimRatio, 0)
    dyBr = np.random.randint(-rows * maxDimRatio, 0)
    dxBl = np.random.randint(0, cols * maxDimRatio)
    dyBl = np.random.randint(-rows * maxDimRatio, 0)
    
    pts1 = np.float32([[0, 0], [cols - 1, 0], [cols - 1, rows - 1], [0, rows - 1]])
    pts2 = np.float32([[0+dxLt, 0+dyLt], [cols-1+dxRt, 0+dyRt], [cols-1+dxBr, rows-1+dyBr], [0+dxBl, rows-1+dyBl]])
    M    = cv2.getPerspectiveTransform(pts1, pts2)
    try:
        Minv = np.linalg.inv(M)
    except:
        return image

    #print('Checking inversion correctness: {}'.format(M.dot(Minv)))
    transformed_image = cv2.warpPerspective(image, M, (cols, rows))
    recovered_image   = cv2.warpPerspective(transformed_image, Minv, (cols, rows))

    return transformed_image, recovered_image
    
def randomPerspectiveTransform(image, minIters=1, maxIters=2):
    iters = np.random.randint(minIters, maxIters)
    for idx in range(iters):
        transformed_image, recovered_image = randomPerspectiveTransformCore(image, maxDimRatio=1./10)
        image                              = recovered_image.copy()    
    return image
def eval_depth(model,global_step):
    torch.cuda.empty_cache()
    save_root=  "./2/" 
    test_path = "./2/"
    for  file in os.listdir(test_path):
        print(file)
        # try:
        # if(".jpg"  in file) :
        if(".bmp"  in file and "IMG20230101103640"  not in file) :
            image = cv2.imread(test_path+file)
            # image = cv2.imread("langan0_left.jpg")
            # in_depth=cv2.imread("langan0_disp.jpg",0)
            # in_depth=cv2.imread(test_path+file.replace(".jpg","_in.png"),0)
            # in_depth=cv2.imread(test_path+file.replace("_left.png","_depth.png"))
            # in_depth=randomPerspectiveTransform(in_depth)[:,:,0]
            in_depth=255-cv2.imread(test_path+file.replace("_bokeh_in_1024x768.bmp","_depthWarp_1024x768.png"),0)
            h,w,c=image.shape
            # image=cv2.resize(image,(480,640))
            # in_depth=cv2.resize(in_depth,(480,640))
            to_tensor = transforms.ToTensor()
            image = to_tensor(image).cuda()
            image = image.unsqueeze(0)
            in_depth = to_tensor(in_depth).cuda()
            in_depth = in_depth.unsqueeze(0)
            image=torch.cat([image,in_depth],dim=1)*255
            with torch.no_grad():
                depth = model(image)
                depth = depth[0,0,:,:]
                # depth=(depth-depth.min())/(depth.max()-depth.min())
                depth=depth
                depth = depth.cpu().numpy()
                depth=np.clip(depth,0,255).astype(np.uint8)
                # depth=cv2.resize(depth,(w,h))
            # depth=cv2.resize(depth,(w,h))
            # cv2.imwrite("kkk.png",depth)
            # cv2.imwrite(test_path+file.replace("_left.png","_refine.png"),depth)
            cv2.imwrite(test_path+file.replace("_bokeh_in_1024x768.bmp","_refine.png"),255-depth)
model=build_efficientnet_lite('efficientnet_lite3', 1000).to(device)
ckpt = torch.load("./model/depth_refine_data_544true_finetune/24-0.7399.pth")
oldCkpt = ckpt
newCkpt = {}
pretrained_state_dict =oldCkpt
model_state_dict=model.state_dict()
for name, param in pretrained_state_dict.items():
    name=name.replace('module.','')
    if name in model_state_dict and model_state_dict[name].shape == param.shape:
        model_state_dict[name] = param
    else:
        print(f"skip loading {name}")
model.load_state_dict(model_state_dict, strict=True)
model.eval()
eval_depth(model, 0)
请把这个改成处理单张图像的代码