Cpt.Captain/CREStereo-pytorch-nxt
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

dataset.py: cleanup

Browse Source
This commit is contained in:
Nils Koch 2022-06-02 15:54:45 +02:00
parent 6ba2cd9e5d
commit 11959eef61
1 changed files with 3 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
dataset.py
View File

@ -86,7 +86,7 @@ class Augmentor:
if len(left_img.shape) == 2:
left_img.shape += 1,
# left_img = cv2.merge([left_img, left_img, left_img])
right_img = cv2.resize(
right_img,
None,
@ -172,7 +172,7 @@ class Augmentor:
if len(left_img.shape) == 2:
left_img = cv2.merge([left_img, left_img, left_img])
return left_img, right_img, left_disp, disp_mask
@ -267,26 +267,14 @@ class CTDDataset(Dataset):
def __getitem__(self, index):
# find path
left_path = self.imgs[index]
prefix = left_path[: left_path.rfind("_")]
# right_path = prefix + "_right.jpg"
left_disp_path = left_path.replace('im', 'disp')
# right_disp_path = prefix + "_right.disp.png"
# read img, disp
# left_img = cv2.imread(left_path, cv2.IMREAD_COLOR)
# right_img = cv2.imread(right_path, cv2.IMREAD_COLOR)
# left_img = np.load(left_path).transpose(1,2,0)
left_img = np.load(left_path)
# left_img = cv2.cvtColor(left_img, cv2.COLOR_GRAY2RGB)
# FIXME DO WE NEED THIS? OTHERWISE IT's PRETTY DARK
left_img = cv2.normalize(left_img, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
# left_img = cv2.merge([left_img, left_img, left_img]).reshape((3, 480, 640))
left_img = cv2.merge([left_img, left_img, left_img]).reshape((480, 640, 3))
right_img = self.pattern
# right_img = cv2.merge([right_img, right_img, right_img]).reshape((3, 480, 640))
left_disp = self.get_disp(left_disp_path)
# right_disp = self.get_disp(right_disp_path)
if False: # self.rng.binomial(1, 0.5):
left_img, right_img = np.fliplr(right_img), np.fliplr(left_img)
@ -308,12 +296,7 @@ class CTDDataset(Dataset):
left_img, right_img, left_disp
)
# left_img = left_img.transpose(2, 0, 1).astype("uint8")
# left_img = left_img.transpose(2, 0, 1).astype("float32")
# left_img = left_img.astype("float32")
right_img = right_img.transpose(2, 0, 1).astype("uint8")
# right_img = right_img.astype("float32")
# print('post_augment', left_img)
right_img = right_img.transpose((2, 0, 1)).astype("uint8")
return {
"left": left_img,