49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import os
|
|
import numpy as np
|
|
import torch
|
|
import torch.utils.data as data
|
|
import torch.nn.functional as F
|
|
import logging
|
|
import os
|
|
import re
|
|
import copy
|
|
import math
|
|
import random
|
|
from pathlib import Path
|
|
from glob import glob
|
|
import os.path as osp
|
|
|
|
from core.utils import frame_utils
|
|
|
|
class CREStereo():
|
|
def __init__(self, aug_params=None, root='/data/CREStereo'):
|
|
self.root = root
|
|
assert os.path.exists(root)
|
|
|
|
disp_list = self.selector('_left.disp.png')
|
|
image1_list = self.selector('_left.jpg')
|
|
image2_list = self.selector('_right.jpg')
|
|
|
|
assert len(image1_list) == len(image2_list) == len(disp_list) > 0
|
|
for img1, img2, disp in zip(image1_list, image2_list, disp_list):
|
|
# if random.randint(1, 20000) != 1:
|
|
# continue
|
|
self.image_list += [[img1, img2]]
|
|
self.disparity_list += [disp]
|
|
|
|
def selector(self, suffix):
|
|
files = list(glob(os.path.join(self.root, f"hole/*{suffix}")))
|
|
files += list(glob(os.path.join(self.root, f"shapenet/*{suffix}")))
|
|
files += list(glob(os.path.join(self.root, f"tree/*{suffix}")))
|
|
files += list(glob(os.path.join(self.root, f"reflective/*{suffix}")))
|
|
return sorted(files)
|
|
|
|
def disparity_distribution(self):
|
|
disp_lists = self.selector('_left.disp.png')
|
|
disparities = []
|
|
for filename in disp_lists:
|
|
disp_gt = frame_utils.readDispCREStereo(filename)
|
|
print(disp_gt.shape)
|
|
|
|
|
|
CREStereo.disparity_distribution() |