113 lines
3.7 KiB
Python
113 lines
3.7 KiB
Python
import os
|
|
import sys
|
|
|
|
sys.path.append("..")
|
|
|
|
import numpy as np
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import random
|
|
from pathlib import Path
|
|
from glob import glob
|
|
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')
|
|
|
|
# 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))
|
|
|
|
|
|
|
|
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 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 = {}
|
|
for filename in disp_lists:
|
|
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()
|
|
|
|
|
|
c = CREStereo()
|
|
c.divide(10000)
|