diff --git a/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py b/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py index 21ea954..f65f2a6 100644 --- a/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py +++ b/IGEV-Stereo/dataset_utils/create_crestereo_subsets.py @@ -1,35 +1,76 @@ import os +import sys + +sys.path.append("..") + 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 shutil import random from pathlib import Path from glob import glob -import os.path as osp +import matplotlib.pyplot as plt from core.utils import frame_utils + +def unique(lst): + return dict(zip(*np.unique(lst, return_counts=True))) + + +def ensure_path_exists(path): + if not os.path.exists(path): + os.makedirs(path) + + 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') + # 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 get_path_info(self, path): + position, filename = os.path.split(path) + root, sub_folder = os.path.split(position) + return root, sub_folder, filename + + def get_new_file(self, path): + root, sub_folder, filename = self.get_path_info(path) + return os.path.join(root, 'subset', sub_folder, filename) + + def divide(self, num): + ensure_path_exists(os.path.join(self.root, 'subset')) + for sub_folder in ['tree', 'shapenet', 'reflective', 'hole']: + ensure_path_exists(os.path.join(self.root, 'subset', sub_folder)) + disp1_list = self.single_folder_selector(sub_folder, '_left.disp.png') + disp2_list = self.single_folder_selector(sub_folder, '_right.disp.png') + image1_list = self.single_folder_selector(sub_folder, '_left.jpg') + image2_list = self.single_folder_selector(sub_folder, '_right.jpg') + + assert len(image1_list) == len(image2_list) == len(disp1_list) == len(disp2_list) > 0 + lists = [] + for img1, img2, disp1, disp2 in zip(image1_list, image2_list, disp1_list, disp2_list): + lists += [[img1, img2, disp1, disp2]] + subset = random.sample(lists, num) + + for s in subset: + for element in s: + print(element) + print(self.get_new_file(element)) + shutil.copy(element, self.get_new_file(element)) + - 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}"))) @@ -38,12 +79,34 @@ class CREStereo(): files += list(glob(os.path.join(self.root, f"reflective/*{suffix}"))) return sorted(files) + def single_folder_selector(self, sub_folder, suffix): + return sorted(list(glob(os.path.join(self.root, f"{sub_folder}/*{suffix}")))) + def disparity_distribution(self): disp_lists = self.selector('_left.disp.png') - disparities = [] + disparities = {} for filename in disp_lists: - disp_gt = frame_utils.readDispCREStereo(filename) - print(disp_gt.shape) + print(filename) + disp_gt, _ = frame_utils.readDispCREStereo(filename) + [rows, cols] = disp_gt.shape + disp_gt = (disp_gt * 32).astype(int) + + cnt = unique(disp_gt) + for i in cnt: + if i in disparities: + disparities[i] += cnt[i] + else: + disparities[i] = cnt[i] + + x = [] + y = [] + for key in disparities.keys(): + x.append(key / 32) + y.append(disparities[key]) + + plt.scatter(x, y) + plt.show() -CREStereo.disparity_distribution() \ No newline at end of file +c = CREStereo() +c.divide(10000)