347 lines
16 KiB
Python
347 lines
16 KiB
Python
import os
|
|
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
|
from typing import IO
|
|
import zipfile
|
|
import json
|
|
import io
|
|
from typing import *
|
|
from pathlib import Path
|
|
import re
|
|
|
|
import numpy as np
|
|
import cv2
|
|
|
|
from .tools import timeit
|
|
|
|
|
|
LEGACY_SEGFORMER_CLASSES = [
|
|
'wall', 'building', 'sky', 'floor', 'tree', 'ceiling', 'road', 'bed ',
|
|
'windowpane', 'grass', 'cabinet', 'sidewalk', 'person', 'earth',
|
|
'door', 'table', 'mountain', 'plant', 'curtain', 'chair', 'car',
|
|
'water', 'painting', 'sofa', 'shelf', 'house', 'sea', 'mirror', 'rug',
|
|
'field', 'armchair', 'seat', 'fence', 'desk', 'rock', 'wardrobe',
|
|
'lamp', 'bathtub', 'railing', 'cushion', 'base', 'box', 'column',
|
|
'signboard', 'chest of drawers', 'counter', 'sand', 'sink',
|
|
'skyscraper', 'fireplace', 'refrigerator', 'grandstand', 'path',
|
|
'stairs', 'runway', 'case', 'pool table', 'pillow', 'screen door',
|
|
'stairway', 'river', 'bridge', 'bookcase', 'blind', 'coffee table',
|
|
'toilet', 'flower', 'book', 'hill', 'bench', 'countertop', 'stove',
|
|
'palm', 'kitchen island', 'computer', 'swivel chair', 'boat', 'bar',
|
|
'arcade machine', 'hovel', 'bus', 'towel', 'light', 'truck', 'tower',
|
|
'chandelier', 'awning', 'streetlight', 'booth', 'television receiver',
|
|
'airplane', 'dirt track', 'apparel', 'pole', 'land', 'bannister',
|
|
'escalator', 'ottoman', 'bottle', 'buffet', 'poster', 'stage', 'van',
|
|
'ship', 'fountain', 'conveyer belt', 'canopy', 'washer', 'plaything',
|
|
'swimming pool', 'stool', 'barrel', 'basket', 'waterfall', 'tent',
|
|
'bag', 'minibike', 'cradle', 'oven', 'ball', 'food', 'step', 'tank',
|
|
'trade name', 'microwave', 'pot', 'animal', 'bicycle', 'lake',
|
|
'dishwasher', 'screen', 'blanket', 'sculpture', 'hood', 'sconce',
|
|
'vase', 'traffic light', 'tray', 'ashcan', 'fan', 'pier', 'crt screen',
|
|
'plate', 'monitor', 'bulletin board', 'shower', 'radiator', 'glass',
|
|
'clock', 'flag'
|
|
]
|
|
LEGACY_SEGFORMER_LABELS = {k: i for i, k in enumerate(LEGACY_SEGFORMER_CLASSES)}
|
|
|
|
|
|
def write_rgbd_zip(
|
|
file: Union[IO, os.PathLike],
|
|
image: Union[np.ndarray, bytes],
|
|
depth: Union[np.ndarray, bytes], mask: Union[np.ndarray, bytes],
|
|
segmentation_mask: Union[np.ndarray, bytes] = None, segmentation_labels: Union[Dict[str, int], bytes] = None,
|
|
intrinsics: np.ndarray = None,
|
|
normal: np.ndarray = None, normal_mask: np.ndarray = None,
|
|
meta: Union[Dict[str, Any], bytes] = None,
|
|
*, image_quality: int = 95, depth_type: Literal['linear', 'log', 'disparity'] = 'linear', depth_format: Literal['png', 'exr'] = 'png', depth_max_dynamic_range: float = 1e4, png_compression: int = 7
|
|
):
|
|
"""
|
|
Write RGBD data as zip archive containing the image, depth, mask, segmentation_mask, and meta data.
|
|
In the zip file there will be:
|
|
- `meta.json`: The meta data as a JSON file.
|
|
- `image.jpg`: The RGB image as a JPEG file.
|
|
- `depth.png/exr`: The depth map as a PNG or EXR file, depending on the `depth_type`.
|
|
- `mask.png` (optional): The mask as a uint8 PNG file.
|
|
- `segmentation_mask.png` (optional): The segformer mask as a uint8/uint16 PNG file.
|
|
|
|
You can provided those data as np.ndarray or bytes. If you provide them as np.ndarray, they will be properly processed and encoded.
|
|
If you provide them as bytes, they will be written as is, assuming they are already encoded.
|
|
"""
|
|
if meta is None:
|
|
meta = {}
|
|
elif isinstance(meta, bytes):
|
|
meta = json.loads(meta.decode())
|
|
|
|
if isinstance(image, bytes):
|
|
image_bytes = image
|
|
elif isinstance(image, np.ndarray):
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
image_bytes = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, image_quality])[1].tobytes()
|
|
|
|
if isinstance(depth, bytes):
|
|
depth_bytes = depth
|
|
elif isinstance(depth, np.ndarray):
|
|
meta['depth_type'] = depth_type
|
|
if depth_type == 'linear':
|
|
if depth.dtype == np.float16:
|
|
depth_format = 'exr'
|
|
depth_bytes = cv2.imencode('.exr', depth.astype(np.float32), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])[1].tobytes()
|
|
elif np.issubdtype(depth.dtype, np.floating):
|
|
depth_format = 'exr'
|
|
depth_bytes = cv2.imencode('.exr', depth.astype(np.float32), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])[1].tobytes()
|
|
elif depth.dtype in [np.uint8, np.uint16]:
|
|
depth_format = 'png'
|
|
depth_bytes = cv2.imencode('.png', depth, [cv2.IMWRITE_PNG_COMPRESSION, png_compression])[1].tobytes()
|
|
elif depth_type == 'log':
|
|
depth_format = 'png'
|
|
depth = depth.astype(np.float32)
|
|
near = max(depth[mask].min(), 1e-3)
|
|
far = min(depth[mask].max(), near * depth_max_dynamic_range)
|
|
depth = ((np.log(depth.clip(near, far) / near) / np.log(far / near)).clip(0, 1) * 65535).astype(np.uint16)
|
|
depth_bytes = cv2.imencode('.png', depth, [cv2.IMWRITE_PNG_COMPRESSION, png_compression])[1].tobytes()
|
|
meta['depth_near'] = float(near)
|
|
meta['depth_far'] = float(far)
|
|
elif depth_type == 'disparity':
|
|
depth_format = 'png'
|
|
depth = depth.astype(np.float32)
|
|
depth = 1 / (depth + 1e-12)
|
|
depth = (depth / depth[mask].max()).clip(0, 1)
|
|
if np.unique(depth) < 200:
|
|
depth = (depth * 255).astype(np.uint8)
|
|
else:
|
|
depth = (depth * 65535).astype(np.uint16)
|
|
depth_bytes = cv2.imencode('.png', depth, [cv2.IMWRITE_PNG_COMPRESSION, png_compression])[1].tobytes()
|
|
|
|
if isinstance(mask, bytes):
|
|
mask_bytes = mask
|
|
elif isinstance(mask, np.ndarray):
|
|
mask_bytes = cv2.imencode('.png', mask.astype(np.uint8) * 255)[1].tobytes()
|
|
|
|
if segmentation_mask is not None:
|
|
if isinstance(segmentation_mask, bytes):
|
|
segmentation_mask_bytes = segmentation_mask
|
|
else:
|
|
segmentation_mask_bytes = cv2.imencode('.png', segmentation_mask)[1].tobytes()
|
|
assert segmentation_labels is not None, "You provided a segmentation mask, but not the corresponding labels."
|
|
if isinstance(segmentation_labels, bytes):
|
|
segmentation_labels = json.loads(segmentation_labels)
|
|
meta['segmentation_labels'] = segmentation_labels
|
|
|
|
if intrinsics is not None:
|
|
meta['intrinsics'] = intrinsics.tolist()
|
|
|
|
if normal is not None:
|
|
if isinstance(normal, bytes):
|
|
normal_bytes = normal
|
|
elif isinstance(normal, np.ndarray):
|
|
normal = ((normal * [0.5, -0.5, -0.5] + 0.5).clip(0, 1) * 65535).astype(np.uint16)
|
|
normal = cv2.cvtColor(normal, cv2.COLOR_RGB2BGR)
|
|
normal_bytes = cv2.imencode('.png', normal, [cv2.IMWRITE_PNG_COMPRESSION, png_compression])[1].tobytes()
|
|
if normal_mask is None:
|
|
normal_mask = np.ones(image.shape[:2], dtype=bool)
|
|
normal_mask_bytes = cv2.imencode('.png', normal_mask.astype(np.uint8) * 255)[1].tobytes()
|
|
|
|
meta_bytes = meta if isinstance(meta, bytes) else json.dumps(meta).encode()
|
|
|
|
with zipfile.ZipFile(file, 'w') as z:
|
|
z.writestr('meta.json', meta_bytes)
|
|
z.writestr('image.jpg', image_bytes)
|
|
z.writestr(f'depth.{depth_format}', depth_bytes)
|
|
z.writestr('mask.png', mask_bytes)
|
|
if segmentation_mask is not None:
|
|
z.writestr('segmentation_mask.png', segmentation_mask_bytes)
|
|
if normal is not None:
|
|
z.writestr('normal.png', normal_bytes)
|
|
z.writestr('normal_mask.png', normal_mask_bytes)
|
|
|
|
|
|
def read_rgbd_zip(file: Union[str, Path, IO], return_bytes: bool = False) -> Dict[str, Union[np.ndarray, Dict[str, Any], bytes]]:
|
|
"""
|
|
Read an RGBD zip file and return the image, depth, mask, segmentation_mask, intrinsics, and meta data.
|
|
|
|
### Parameters:
|
|
- `file: Union[str, Path, IO]`
|
|
The file path or file object to read from.
|
|
- `return_bytes: bool = False`
|
|
If True, return the image, depth, mask, and segmentation_mask as raw bytes.
|
|
|
|
### Returns:
|
|
- `Tuple[Dict[str, Union[np.ndarray, Dict[str, Any]]], Dict[str, bytes]]`
|
|
A dictionary containing: (If missing, the value will be None; if return_bytes is True, the value will be bytes)
|
|
- `image`: RGB numpy.ndarray of shape (H, W, 3).
|
|
- `depth`: float32 numpy.ndarray of shape (H, W).
|
|
- `mask`: bool numpy.ndarray of shape (H, W).
|
|
- `segformer_mask`: uint8 numpy.ndarray of shape (H, W).
|
|
- `intrinsics`: float32 numpy.ndarray of shape (3, 3).
|
|
- `meta`: Dict[str, Any].
|
|
"""
|
|
# Load & extract archive
|
|
with zipfile.ZipFile(file, 'r') as z:
|
|
meta = z.read('meta.json')
|
|
if not return_bytes:
|
|
meta = json.loads(z.read('meta.json'))
|
|
|
|
image = z.read('image.jpg')
|
|
if not return_bytes:
|
|
image = cv2.imdecode(np.frombuffer(z.read('image.jpg'), np.uint8), cv2.IMREAD_COLOR)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
depth_name = next(s for s in z.namelist() if s.startswith('depth'))
|
|
depth = z.read(depth_name)
|
|
if not return_bytes:
|
|
depth = cv2.imdecode(np.frombuffer(z.read(depth_name), np.uint8), cv2.IMREAD_UNCHANGED)
|
|
|
|
if 'mask.png' in z.namelist():
|
|
mask = z.read('mask.png')
|
|
if not return_bytes:
|
|
mask = cv2.imdecode(np.frombuffer(z.read('mask.png'), np.uint8), cv2.IMREAD_UNCHANGED) > 0
|
|
else:
|
|
mask = None
|
|
|
|
if 'segformer_mask.png' in z.namelist():
|
|
# NOTE: Legacy support for segformer_mask.png
|
|
segmentation_mask = z.read('segformer_mask.png')
|
|
segmentation_labels = None
|
|
if not return_bytes:
|
|
segmentation_mask = cv2.imdecode(np.frombuffer(segmentation_mask, np.uint8), cv2.IMREAD_UNCHANGED)
|
|
segmentation_labels = LEGACY_SEGFORMER_LABELS
|
|
elif 'segmentation_mask.png' in z.namelist():
|
|
segmentation_mask = z.read('segmentation_mask.png')
|
|
segmentation_labels = None
|
|
if not return_bytes:
|
|
segmentation_mask = cv2.imdecode(np.frombuffer(segmentation_mask, np.uint8), cv2.IMREAD_UNCHANGED)
|
|
segmentation_labels = meta['segmentation_labels']
|
|
else:
|
|
segmentation_mask = None
|
|
segmentation_labels = None
|
|
|
|
if 'normal.png' in z.namelist():
|
|
normal = z.read('normal.png')
|
|
if not return_bytes:
|
|
normal = cv2.imdecode(np.frombuffer(z.read('normal.png'), np.uint8), cv2.IMREAD_UNCHANGED)
|
|
normal = cv2.cvtColor(normal, cv2.COLOR_BGR2RGB)
|
|
normal = (normal.astype(np.float32) / 65535 - 0.5) * [2.0, -2.0, -2.0]
|
|
normal = normal / np.linalg.norm(normal, axis=-1, keepdims=True)
|
|
|
|
if 'normal_mask.png' in z.namelist():
|
|
normal_mask = z.read('normal_mask.png')
|
|
normal_mask = cv2.imdecode(np.frombuffer(normal_mask, np.uint8), cv2.IMREAD_UNCHANGED) > 0
|
|
else:
|
|
normal_mask = np.ones(image.shape[:2], dtype=bool)
|
|
else:
|
|
normal, normal_mask = None, None
|
|
|
|
# recover linear depth
|
|
if not return_bytes:
|
|
if mask is None:
|
|
mask = np.ones(image.shape[:2], dtype=bool)
|
|
if meta['depth_type'] == 'linear':
|
|
depth = depth.astype(np.float32)
|
|
mask = mask & (depth > 0)
|
|
elif meta['depth_type'] == 'log':
|
|
near, far = meta['depth_near'], meta['depth_far']
|
|
if depth.dtype == np.uint16:
|
|
depth = depth.astype(np.float32) / 65535
|
|
elif depth.dtype == np.uint8:
|
|
depth = depth.astype(np.float32) / 255
|
|
depth = near ** (1 - depth) * far ** depth
|
|
mask = mask & ~np.isnan(depth)
|
|
elif meta['depth_type'] == 'disparity':
|
|
mask = mask & (depth > 0)
|
|
if depth.dtype == np.uint16:
|
|
depth = depth.astype(np.float32) / 65535
|
|
elif depth.dtype == np.uint8:
|
|
depth = depth.astype(np.float32) / 255
|
|
depth = 1 / (depth + 1e-12)
|
|
|
|
# intrinsics
|
|
if not return_bytes and 'intrinsics' in meta:
|
|
intrinsics = np.array(meta['intrinsics'], dtype=np.float32)
|
|
else:
|
|
intrinsics = None
|
|
|
|
# depth unit
|
|
if not return_bytes and 'depth_unit' in meta:
|
|
depth_unit_str = meta['depth_unit']
|
|
if r := re.match(r'([\d.]*)(\w*)', depth_unit_str):
|
|
digits, unit = r.groups()
|
|
depth_unit = float(digits or 1) * {'m': 1, 'cm': 0.01, 'mm': 0.001}[unit]
|
|
else:
|
|
depth_unit = None
|
|
else:
|
|
depth_unit = None
|
|
|
|
return_dict = {
|
|
'image': image,
|
|
'depth': depth,
|
|
'mask': mask,
|
|
'segmentation_mask': segmentation_mask,
|
|
'segmentation_labels': segmentation_labels,
|
|
'normal': normal,
|
|
'normal_mask': normal_mask,
|
|
'intrinsics': intrinsics,
|
|
'depth_unit': depth_unit,
|
|
'meta': meta,
|
|
}
|
|
return_dict = {k: v for k, v in return_dict.items() if v is not None}
|
|
|
|
return return_dict
|
|
|
|
def write_rgbxyz(file: Union[IO, Path], image: np.ndarray, points: np.ndarray, mask: np.ndarray = None, image_quality: int = 95):
|
|
if isinstance(image, bytes):
|
|
image_bytes = image
|
|
elif isinstance(image, np.ndarray):
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
image_bytes = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, image_quality])[1].tobytes()
|
|
|
|
if isinstance(points, bytes):
|
|
points_bytes = points
|
|
elif isinstance(points, np.ndarray):
|
|
points_bytes = cv2.imencode('.exr', points.astype(np.float32), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT])[1].tobytes()
|
|
|
|
if mask is None:
|
|
mask = np.ones(image.shape[:2], dtype=bool)
|
|
if isinstance(mask, bytes):
|
|
mask_bytes = mask
|
|
elif isinstance(mask, np.ndarray):
|
|
mask_bytes = cv2.imencode('.png', mask.astype(np.uint8) * 255)[1].tobytes()
|
|
|
|
is_archive = hasattr(file, 'write') or Path(file).suffix == '.zip'
|
|
if is_archive:
|
|
with zipfile.ZipFile(file, 'w') as z:
|
|
z.writestr('image.jpg', image_bytes)
|
|
z.writestr('points.exr', points_bytes)
|
|
if mask is not None:
|
|
z.writestr('mask.png', mask_bytes)
|
|
else:
|
|
file = Path(file)
|
|
file.mkdir(parents=True, exist_ok=True)
|
|
with open(file / 'image.jpg', 'wb') as f:
|
|
f.write(image_bytes)
|
|
with open(file / 'points.exr', 'wb') as f:
|
|
f.write(points_bytes)
|
|
if mask is not None:
|
|
with open(file / 'mask.png', 'wb') as f:
|
|
f.write(mask_bytes)
|
|
|
|
|
|
def read_rgbxyz(file: Union[IO, str, Path]) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
|
|
is_archive = hasattr(file, 'read') or Path(file).suffix == '.zip'
|
|
if is_archive:
|
|
with zipfile.ZipFile(file, 'r') as z:
|
|
image = cv2.cvtColor(cv2.imdecode(np.frombuffer(z.read('image.jpg'), np.uint8), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
|
|
points = cv2.cvtColor(cv2.imdecode(np.frombuffer(z.read('points.exr'), np.uint8), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
|
|
if 'mask.png' in z.namelist():
|
|
mask = cv2.imdecode(np.frombuffer(z.read('mask.png'), np.uint8), cv2.IMREAD_GRAYSCALE) > 0
|
|
else:
|
|
mask = np.ones(image.shape[:2], dtype=bool)
|
|
else:
|
|
file = Path(file)
|
|
file.mkdir(parents=True, exist_ok=True)
|
|
image = cv2.cvtColor(cv2.imread(str(file / 'image.jpg'), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
|
|
points = cv2.cvtColor(cv2.imread(str(file / 'points.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB)
|
|
if (file /'mask.png').exists():
|
|
mask = cv2.imread(str(file / 'mask.png'), cv2.IMREAD_GRAYSCALE) > 0
|
|
else:
|
|
mask = np.ones(image.shape[:2], dtype=bool)
|
|
|
|
return image, points, mask
|