Custom phototourism data loader

datasets/phototourism.py

@@ -52,13 +52,18 @@ def read_meta(self):
         if self.use_cache:
             with open(os.path.join(self.root_dir, f'cache/img_ids.pkl'), 'rb') as f:
                 self.img_ids = pickle.load(f)
+            with open(os.path.join(self.root_dir, f'cache/img_to_cam_id.pkl'), 'rb') as f:
+                self.image_to_cam = pickle.load(f)
             with open(os.path.join(self.root_dir, f'cache/image_paths.pkl'), 'rb') as f:
                 self.image_paths = pickle.load(f)
         else:
             imdata = read_images_binary(os.path.join(self.root_dir, 'dense/sparse/images.bin'))
             img_path_to_id = {}
+            self.image_to_cam = {} # {id: image id}
+
             for v in imdata.values():
                 img_path_to_id[v.name] = v.id
+                self.image_to_cam[v.id] = v.camera_id
             self.img_ids = []
             self.image_paths = {} # {id: filename}
             for filename in list(self.files['filename']):
@@ -75,15 +80,17 @@ def read_meta(self):
             camdata = read_cameras_binary(os.path.join(self.root_dir, 'dense/sparse/cameras.bin'))
             for id_ in self.img_ids:
                 K = np.zeros((3, 3), dtype=np.float32)
-                cam = camdata[id_]
+
+                cam_id = self.image_to_cam[id_]
+                cam = camdata[cam_id]
                 img_w, img_h = int(cam.params[2]*2), int(cam.params[3]*2)
                 img_w_, img_h_ = img_w//self.img_downscale, img_h//self.img_downscale
                 K[0, 0] = cam.params[0]*img_w_/img_w # fx
                 K[1, 1] = cam.params[1]*img_h_/img_h # fy
                 K[0, 2] = cam.params[2]*img_w_/img_w # cx
                 K[1, 2] = cam.params[3]*img_h_/img_h # cy
                 K[2, 2] = 1
-                self.Ks[id_] = K
+                self.Ks[cam_id] = K
 
         # Step 3: read c2w poses (of the images in tsv file only) and correct the order
         if self.use_cache:
@@ -163,7 +170,7 @@ def read_meta(self):
                     img = img.view(3, -1).permute(1, 0) # (h*w, 3) RGB
                     self.all_rgbs += [img]
 
-                    directions = get_ray_directions(img_h, img_w, self.Ks[id_])
+                    directions = get_ray_directions(img_h, img_w, self.Ks[self.image_to_cam[id_]])
                     rays_o, rays_d = get_rays(directions, c2w)
                     rays_t = id_ * torch.ones(len(rays_o), 1)
 
@@ -220,7 +227,7 @@ def __getitem__(self, idx):
             img = img.view(3, -1).permute(1, 0) # (h*w, 3) RGB
             sample['rgbs'] = img
 
-            directions = get_ray_directions(img_h, img_w, self.Ks[id_])
+            directions = get_ray_directions(img_h, img_w, self.Ks[self.image_to_cam[id_]])
             rays_o, rays_d = get_rays(directions, c2w)
             rays = torch.cat([rays_o, rays_d,
                               self.nears[id_]*torch.ones_like(rays_o[:, :1]),

prepare_phototourism.py

@@ -23,6 +23,8 @@ def get_opts():
     # save img ids
     with open(os.path.join(args.root_dir, f'cache/img_ids.pkl'), 'wb') as f:
         pickle.dump(dataset.img_ids, f, pickle.HIGHEST_PROTOCOL)
+    with open(os.path.join(args.root_dir, f'cache/img_to_cam_id.pkl'), 'wb') as f:
+        pickle.dump(dataset.image_to_cam, f, pickle.HIGHEST_PROTOCOL)
     # save img paths
     with open(os.path.join(args.root_dir, f'cache/image_paths.pkl'), 'wb') as f:
         pickle.dump(dataset.image_paths, f, pickle.HIGHEST_PROTOCOL)

utils/gen_nerf_tsv.py

@@ -0,0 +1,23 @@
+import pandas as pd
+import os
+
+dataset_name = 'nepszinhaz'
+path = f'/mnt/hdd/datasets/budapest/{dataset_name}'
+images = f'{path}/dense/images'
+out_path = f'{path}/{dataset_name}.tsv'
+
+db = []
+
+header = ['filename', 'id', 'split', 'dataset']
+
+for image in os.listdir(images):
+    # todo set proper train/test
+    db.append([image, '-1', 'train', dataset_name])
+
+with open(out_path, 'w') as f:
+    f.write('\t'.join(header))
+    f.write('\n')
+
+    for data in db:
+        f.write('\t'.join(data))
+        f.write('\n')