moge/test/dataloader.py

import os
from typing import *
from pathlib import Path
import math

import numpy as np
import torch
from PIL import Image
import cv2
try:
    import utils3d_moge as utils3d
except ImportError:
    import utils3d
import pipeline

from ..utils.geometry_numpy import focal_to_fov_numpy, norm3d
from ..utils.io import *
from ..utils.tools import timeit


class EvalDataLoaderPipeline:
    def __init__(
        self,
        path: str,
        width: int,
        height: int,
        split: str = '.index.txt',
        depth: str = 'depth.png',
        drop_max_depth: float = 1000.,
        segmentation: str = None,
        normal: str = None,
        local_mask: str = None,
        local_segmentation: str = None,
        max_segments: int = 100,
        min_seg_area: int = 1000,
        depth_unit: str = None,
        has_sharp_boundary = False,
        subset: int = None,
    ):
        num_load_workers = int(os.environ.get('MOGE_NUM_LOAD_WORKERS', 4))
        num_process_workers = int(os.environ.get('MOGE_NUM_PROCESS_WORKERS', 8))
    
        filenames = Path(path).joinpath(split).read_text(encoding='utf-8').splitlines()
        filenames = filenames[::subset]
        self.width = width
        self.height = height
        self.drop_max_depth = drop_max_depth
        self.path = Path(path)
        self.filenames = filenames
        self.depth_file = depth
        self.segmentation_file = segmentation
        self.normal_file = normal
        self.local_mask_file = local_mask
        self.local_segmentation_file = local_segmentation
        self.max_segments = max_segments
        self.min_seg_area = min_seg_area
        self.depth_unit = depth_unit
        self.has_sharp_boundary = has_sharp_boundary

        self.rng = np.random.default_rng(seed=0)
        
        self.pipeline = pipeline.Sequential([
            self._generator,
            pipeline.Parallel([self._load_instance] * num_load_workers),
            pipeline.Parallel([self._process_instance] * num_process_workers),
            pipeline.Buffer(4)
        ])

    def __len__(self):
        return math.ceil(len(self.filenames)) 

    def _generator(self):
        for idx in range(len(self)):
            yield idx
    
    def _load_instance(self, idx):
        if idx >= len(self.filenames):
            return None
        
        path = self.path.joinpath(self.filenames[idx])

        instance = {
            'filename': self.filenames[idx],
            'width': self.width,
            'height': self.height,
        }
        instance['image'] = read_image(Path(path, 'image.jpg'))

        depth = read_depth(Path(path, self.depth_file))  # ignore depth unit from depth file, use config instead
        instance.update({
            'depth': np.nan_to_num(depth, nan=1, posinf=1, neginf=1),
            'depth_mask': np.isfinite(depth),
            'depth_mask_inf': np.isinf(depth),
        })

        if self.segmentation_file is not None:
            segmentation_mask, segmentation_labels = read_segmentation(Path(path, self.segmentation_file))
            if segmentation_labels is None:
                # SAM-style id maps (sam_segmentation*.png) carry no PNG `labels`
                # metadata. Synthesize one entry per non-zero id so the downstream
                # label filtering in _process_instance still works.
                segmentation_labels = {f'{int(i):05d}': int(i) for i in np.unique(segmentation_mask) if i > 0}
            instance.update({
                'segmentation_mask': segmentation_mask,
                'segmentation_labels': segmentation_labels,
            })

        meta = read_json(Path(path, 'meta.json'))
        instance['intrinsics'] = np.array(meta['intrinsics'], dtype=np.float32)

        if self.normal_file is not None:
            normal = read_normal(Path(path, self.normal_file))
            # `read_normal` marks invalid pixels with NaN, so derive the mask
            # before scrubbing them.
            normal_mask = np.isfinite(normal).all(axis=-1)
            normal = np.nan_to_num(normal, nan=0, posinf=0, neginf=0)
            instance.update({
                'normal': normal,
                'normal_mask': normal_mask,
            })

        if self.local_mask_file is not None:
            local_mask_path = Path(path, self.local_mask_file)
            if local_mask_path.exists():
                local_mask = cv2.imread(str(local_mask_path), cv2.IMREAD_GRAYSCALE)
                instance['local_mask'] = local_mask > 0

                if self.local_segmentation_file is not None:
                    hf_seg_path = Path(path, self.local_segmentation_file)
                    if hf_seg_path.exists():
                        hf_seg = cv2.imread(str(hf_seg_path), cv2.IMREAD_UNCHANGED)
                        if hf_seg.shape[:2] != instance['local_mask'].shape:
                            hf_seg = cv2.resize(
                                hf_seg,
                                (instance['local_mask'].shape[1], instance['local_mask'].shape[0]),
                                interpolation=cv2.INTER_NEAREST,
                            )
                        instance['local_segmentation'] = hf_seg.astype(np.int32)

        return instance

    def _process_instance(self, instance: dict):
        if instance is None:
            return None
        
        image, depth, depth_mask, intrinsics = instance['image'], instance['depth'], instance['depth_mask'], instance['intrinsics']
        segmentation_mask, segmentation_labels = instance.get('segmentation_mask', None), instance.get('segmentation_labels', None)
        normal, normal_mask = instance.get('normal', None), instance.get('normal_mask', None)

        raw_height, raw_width = image.shape[:2]
        raw_horizontal, raw_vertical = abs(1.0 / intrinsics[0, 0]), abs(1.0 / intrinsics[1, 1])
        raw_pixel_w, raw_pixel_h = raw_horizontal / raw_width, raw_vertical / raw_height
        tgt_width, tgt_height = instance['width'], instance['height']
        tgt_aspect = tgt_width / tgt_height

        # set expected target view field
        tgt_horizontal = min(raw_horizontal, raw_vertical * tgt_aspect)
        tgt_vertical = tgt_horizontal / tgt_aspect

        # set target view direction
        cu, cv = 0.5, 0.5
        direction = utils3d.np.unproject_cv(np.array([[cu, cv]], dtype=np.float32), np.array([1.0], dtype=np.float32), intrinsics=intrinsics)[0]
        R = utils3d.np.rotation_matrix_from_vectors(direction, np.array([0, 0, 1], dtype=np.float32))

        # restrict target view field within the raw view
        corners = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32)
        corners = np.concatenate([corners, np.ones((4, 1), dtype=np.float32)], axis=1) @ (np.linalg.inv(intrinsics).T @ R.T)   # corners in viewport's camera plane
        corners = corners[:, :2] / corners[:, 2:3]

        warp_horizontal, warp_vertical = abs(1.0 / intrinsics[0, 0]), abs(1.0 / intrinsics[1, 1])
        for i in range(4):
            intersection, _ = utils3d.np.ray_intersection(
                np.array([0., 0.]), np.array([[tgt_aspect, 1.0], [tgt_aspect, -1.0]]),
                corners[i - 1], corners[i] - corners[i - 1],
            )
            warp_horizontal, warp_vertical = min(warp_horizontal, 2 * np.abs(intersection[:, 0]).min()), min(warp_vertical, 2 * np.abs(intersection[:, 1]).min())
        tgt_horizontal, tgt_vertical = min(tgt_horizontal, warp_horizontal), min(tgt_vertical, warp_vertical)

        # get target view intrinsics
        fx, fy = 1.0 / tgt_horizontal, 1.0 / tgt_vertical
        tgt_intrinsics = utils3d.np.intrinsics_from_focal_center(fx, fy, 0.5, 0.5).astype(np.float32)
        
        # do homogeneous transformation with the rotation and intrinsics
        # 4.1 The image and depth is resized first to approximately the same pixel size as the target image with PIL's antialiasing resampling
        tgt_pixel_w, tgt_pixel_h = tgt_horizontal / tgt_width, tgt_vertical / tgt_height        # (should be exactly the same for x and y axes)
        rescaled_w, rescaled_h = int(raw_width * raw_pixel_w / tgt_pixel_w), int(raw_height * raw_pixel_h / tgt_pixel_h)
        image = np.array(Image.fromarray(image).resize((rescaled_w, rescaled_h), Image.Resampling.LANCZOS))

        depth, depth_mask = utils3d.np.masked_nearest_resize(depth, mask=depth_mask, size=(rescaled_h, rescaled_w))
        distance = norm3d(utils3d.np.depth_map_to_point_map(depth, intrinsics=intrinsics))
        segmentation_mask = cv2.resize(segmentation_mask, (rescaled_w, rescaled_h), interpolation=cv2.INTER_NEAREST) if segmentation_mask is not None else None
        if normal is not None:
            # NOTE: `utils3d.np.masked_nearest_resize` takes size as (H, W) -- the
            # opposite of cv2's (W, H) and of `mask_aware_nearest_resize_numpy`.
            normal, normal_mask = utils3d.np.masked_nearest_resize(normal, mask=normal_mask, size=(rescaled_h, rescaled_w))

        # 4.2 calculate homography warping
        transform = intrinsics @ np.linalg.inv(R) @ np.linalg.inv(tgt_intrinsics)
        uv_tgt = utils3d.np.uv_map(tgt_height, tgt_width)
        pts = np.concatenate([uv_tgt, np.ones((tgt_height, tgt_width, 1), dtype=np.float32)], axis=-1) @ transform.T
        uv_remap = pts[:, :, :2] / (pts[:, :, 2:3] + 1e-12)
        pixel_remap = utils3d.np.uv_to_pixel(uv_remap, (rescaled_h, rescaled_w)).astype(np.float32)
        
        tgt_image = cv2.remap(image, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LINEAR)
        tgt_distance = cv2.remap(distance, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST)
        tgt_ray_length = utils3d.np.unproject_cv(uv_tgt, np.ones_like(uv_tgt[:, :, 0]), intrinsics=tgt_intrinsics)
        tgt_ray_length = (tgt_ray_length[:, :, 0] ** 2 + tgt_ray_length[:, :, 1] ** 2 + tgt_ray_length[:, :, 2] ** 2) ** 0.5
        tgt_depth = tgt_distance / (tgt_ray_length + 1e-12)
        tgt_depth_mask = cv2.remap(depth_mask.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0
        tgt_segmentation_mask = cv2.remap(segmentation_mask, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) if segmentation_mask is not None else None
        tgt_normal = cv2.remap(normal, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LINEAR) @ R.T if normal is not None else None
        tgt_normal_mask = cv2.remap(normal_mask.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0 if normal_mask is not None else None

        # NOTE: cv2.resize / cv2.remap take size as (W, H), unlike the
        # `masked_nearest_resize` call above.
        local_mask = instance.get('local_mask', None)
        if local_mask is not None:
            local_mask = cv2.resize(local_mask.astype(np.uint8), (rescaled_w, rescaled_h), interpolation=cv2.INTER_NEAREST)
            tgt_local_mask = cv2.remap(local_mask, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0
        else:
            tgt_local_mask = None

        local_segmentation = instance.get('local_segmentation', None)
        if local_segmentation is not None:
            # int32 round-trips through resize/remap under INTER_NEAREST.
            local_segmentation = cv2.resize(local_segmentation, (rescaled_w, rescaled_h), interpolation=cv2.INTER_NEAREST)
            tgt_local_segmentation = cv2.remap(local_segmentation, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST)
        else:
            tgt_local_segmentation = None

        # drop depth greater than drop_max_depth
        max_depth = np.nanquantile(np.where(tgt_depth_mask, tgt_depth, np.nan), 0.01) * self.drop_max_depth
        tgt_depth_mask &= tgt_depth <= max_depth
        tgt_depth = np.nan_to_num(tgt_depth, nan=0.0)

        if self.depth_unit is not None:
            tgt_depth *= self.depth_unit
        
        if not np.any(tgt_depth_mask):
            # always make sure that mask is not empty, otherwise the loss calculation will crash
            tgt_depth_mask = np.ones_like(tgt_depth_mask)
            tgt_depth = np.ones_like(tgt_depth)
            instance['label_type'] = 'invalid'
        
        tgt_pts = utils3d.np.unproject_cv(uv_tgt, tgt_depth, intrinsics=tgt_intrinsics)

        # Process segmentation labels
        if segmentation_mask is not None:
            for k in ['undefined', 'unannotated', 'background', 'sky']:
                if k in segmentation_labels:
                    del segmentation_labels[k]
            seg_id2count = dict(zip(*np.unique(tgt_segmentation_mask, return_counts=True)))
            sorted_labels = sorted(segmentation_labels.keys(), key=lambda x: seg_id2count.get(segmentation_labels[x], 0), reverse=True)
            segmentation_labels = {k: segmentation_labels[k] for k in sorted_labels[:self.max_segments] if seg_id2count.get(segmentation_labels[k], 0) >= self.min_seg_area}

        instance.update({
            'image': torch.from_numpy(tgt_image.astype(np.float32) / 255.0).permute(2, 0, 1),
            'depth': torch.from_numpy(tgt_depth).float(),
            'depth_mask': torch.from_numpy(tgt_depth_mask).bool(),
            'intrinsics': torch.from_numpy(tgt_intrinsics).float(),
            'points': torch.from_numpy(tgt_pts).float(),
            'normal': torch.from_numpy(tgt_normal).float() if tgt_normal is not None else None,
            'normal_mask': torch.from_numpy(tgt_normal_mask).bool() if tgt_normal_mask is not None else None,
            'segmentation_mask': torch.from_numpy(tgt_segmentation_mask).long() if tgt_segmentation_mask is not None else None,
            'segmentation_labels': segmentation_labels,
            'local_mask': torch.from_numpy(tgt_local_mask).bool() if tgt_local_mask is not None else None,
            'local_segmentation': torch.from_numpy(tgt_local_segmentation).long() if tgt_local_segmentation is not None else None,
            'is_metric': self.depth_unit is not None,
            'has_sharp_boundary': self.has_sharp_boundary,
        })
        
        instance = {k: v for k, v in instance.items() if v is not None}
        
        return instance

    def start(self):
        self.pipeline.start()
    
    def stop(self):
        self.pipeline.stop()
    
    def __enter__(self):
        self.start()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.stop()

    def get(self):
        return self.pipeline.get()