def loadData()

in varifocal/code/sync_defocus_from_rgbd_fast/train/preprocess.py [0:0]


def loadData(data_path, N, H, W, imH, imW,dp_focal,fov, D, startSceneIDX,endSceneIDX,numEvalScenes,numChannels, depthScale, diopterScale, LOAD_FS, LOAD_RGBD):
    print("preprocessing data...")

    numScenes = 0
    validSceneIdxList = []
    for s in range(startSceneIDX, endSceneIDX+1):
        if os.path.exists("%s/seed%04d/" % (data_path, s)):
           validSceneIdxList.append(s)
           numScenes = numScenes+1

    print("found %d valid scenes" % numScenes)
    if numScenes==0:
       input("error: please correct data path!")

    patch_shift = np.int(np.floor(W/2)) #W/2 # when patch_shift<W, overlapped patches are allowed
    num_patch_per_row = np.int(np.floor((imW-W)/patch_shift+1))
    num_patch_per_col = np.int(np.floor((imH-H)/patch_shift+1))
    num_patch_per_img = num_patch_per_row*num_patch_per_col
    numSamples = numScenes*num_patch_per_img # number of sample patches

    numTrainSamples = (numScenes-numEvalScenes)*num_patch_per_img # number of samples for training
    numEvalSamples = numSamples - numTrainSamples # number of samples for validation during training
    
    focal_stack = np.zeros((numSamples, H, W, numChannels, N), dtype=np.float32) 
    cleanRGB = np.zeros((numSamples, H, W, numChannels), dtype=np.float32)
    depth = np.zeros((numSamples, H, W), dtype=np.float32)
    diopter = np.zeros((numSamples, H, W), dtype=np.float32)

    
    for s in range(0, numScenes): # 1:numScenes
       
       v = validSceneIdxList[s]
       print("load scene %03d" % v)
       
       # load RGBD data
       if LOAD_RGBD:
           # load sharp RGB images
           im = cv2.imread("%s/seed%04d/%d/clean_pass_rgb.exr" % (data_path, v, imW), -1)             
           im, im_max, im_min = normalize_rgb_exr(im)

           for row in range(0, num_patch_per_row):
             for col in range(0, num_patch_per_col):
                   patch = im[row*patch_shift:row*patch_shift+H, col*patch_shift:col*patch_shift+W, :numChannels]
                   idx_patch = row*num_patch_per_col + col
                   idx = s*num_patch_per_img + idx_patch
                   cleanRGB[idx,:,:,:] = patch 
 

           # load depth map
           im = cv2.imread("%s/seed%04d/%d/clean_pass_depth_rgb.exr" % (data_path, v, imW), -1)
           im = rescale_depth_exr(im, depthScale) # use global scale (not per scene)
           for row in range(0, num_patch_per_row):
             for col in range(0, num_patch_per_col):
                patch = im[row*patch_shift:row*patch_shift+H, col*patch_shift:col*patch_shift+W, :]
                idx_patch = row*num_patch_per_col + col
                idx = s*num_patch_per_img+idx_patch
                depth[idx,:,:] = patch[:,:,0]
                diopter[idx,:,:] = 1/(depthScale*depth[idx,:,:])/diopterScale

                   
       # load focal stack images
       if LOAD_FS:
          for i in range(0, len(dp_focal)):
             #print("scene:%d, image:%d" % (s, i))
             im = cv2.imread("%s/seed%04d/%d/frame%04d.exr" % (data_path, v, imW, i), -1)             
             im = rescale_rgb_exr(im, im_max, im_min)
             for row in range(0, num_patch_per_row):
                for col in range(0, num_patch_per_col):
                   patch = im[row*patch_shift:row*patch_shift+H, col*patch_shift:col*patch_shift+W, :numChannels]
                   idx_patch = row*num_patch_per_col + col
                   idx = s*num_patch_per_img+idx_patch
                   focal_stack[idx,:,:,:,i] = patch                                    


    print("data imported")
    return focal_stack, dp_focal, numScenes, numTrainSamples, numEvalSamples, cleanRGB, depth, diopter,validSceneIdxList