datasets/matterport.py

import numpy as np
import cv2
from pathlib import Path
import json
import torch
from torch.utils.data import Dataset

from utils import rotation_matrix_from_quaternion, Augmentor


class Matterport3D(Dataset):
    def __init__(self, data_root, mode='train'):
        data_root = Path(data_root)
        json_path = data_root / 'mp3d_planercnn_json' / f'cached_set_{mode}.json'

        scene_info = {'images': [], 'rotation': [], 'translation': [], 'intrinsics': []}

        with open(json_path) as f:
            split = json.load(f)
        
        for _, data in enumerate(split['data']):
            images = []
            for imgnum in ['0', '1']:
                img_name = data_root / '/'.join(data[imgnum]['file_name'].split('/')[6:])
                images.append(img_name)
            
            rel_rotation = data['rel_pose']['rotation']
            rel_translation = data['rel_pose']['position']
            intrinsic = [
                [517.97, 0, 320],
                [0, 517.97, 240],
                [0, 0, 1]
            ]
            intrinsics = [intrinsic, intrinsic]

            scene_info['images'].append(images)
            scene_info['rotation'].append(rel_rotation)
            scene_info['translation'].append(rel_translation)
            scene_info['intrinsics'].append(intrinsics)

        scene_info['rotation'] = torch.tensor(scene_info['rotation'])
        scene_info['translation'] = torch.tensor(scene_info['translation'])
        scene_info['intrinsics'] = torch.tensor(scene_info['intrinsics'])

        self.scene_info = scene_info
        self.augment = Augmentor(mode=='train')

        self.is_training = mode == 'train'

    def __len__(self):
        return len(self.scene_info['images'])

    def __getitem__(self, idx):
        img_names = self.scene_info['images'][idx]
        rotation = self.scene_info['rotation'][idx]
        translation = self.scene_info['translation'][idx]
        intrinsics = self.scene_info['intrinsics'][idx]

        images = []
        for i in range(2):
            image = cv2.imread(str(img_names[i]))
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image = self.augment(image)
            image = torch.from_numpy(image).permute(2, 0, 1)
            images.append(image)
        images = torch.stack(images)
        images = images.float() / 255.

        rotation = -rotation if rotation[0] < 0 else rotation
        rotation /= rotation.norm(2)
        rotation = rotation_matrix_from_quaternion(rotation[None,])[0]

        rotation = rotation.mT
        translation = -rotation @ translation.unsqueeze(-1)
        translation = translation[:, 0]

        return {
            'images': images,
            'rotation': rotation,
            'translation': translation,
            'intrinsics': intrinsics,
        }


def build_matterport(mode, config):
    return Matterport3D(config.DATASET.DATA_ROOT, mode)