Create create_crestereo_subsets.py but haven't finished

IGEV-Stereo/dataset_utils/create_crestereo_subsets.py

@@ -0,0 +1,49 @@
+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()