Spaces:

heheyas
/

V3D

Sleeping

V3D / sgm /data /json_index_dataset.py

heheyas

init

cfb7702 3 months ago

No virus

42.5 kB

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.

	import copy
	import functools
	import gzip
	import hashlib
	import json
	import logging
	import os
	import random
	import warnings
	from collections import defaultdict
	from itertools import islice
	from pathlib import Path
	from typing import (
	Any,
	ClassVar,
	Dict,
	Iterable,
	List,
	Optional,
	Sequence,
	Tuple,
	Type,
	TYPE_CHECKING,
	Union,
	)

	import numpy as np
	import torch
	from PIL import Image
	from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
	from pytorch3d.io import IO
	from pytorch3d.renderer.camera_utils import join_cameras_as_batch
	from pytorch3d.renderer.cameras import CamerasBase, PerspectiveCameras
	from pytorch3d.structures.pointclouds import Pointclouds
	from tqdm import tqdm

	from pytorch3d.implicitron.dataset import types
	from pytorch3d.implicitron.dataset.dataset_base import DatasetBase, FrameData
	from pytorch3d.implicitron.dataset.utils import is_known_frame_scalar


	logger = logging.getLogger(__name__)


	if TYPE_CHECKING:
	from typing import TypedDict

	class FrameAnnotsEntry(TypedDict):
	subset: Optional[str]
	frame_annotation: types.FrameAnnotation

	else:
	FrameAnnotsEntry = dict


	@registry.register
	class JsonIndexDataset(DatasetBase, ReplaceableBase):
	"""
	A dataset with annotations in json files like the Common Objects in 3D
	(CO3D) dataset.

	Args:
	frame_annotations_file: A zipped json file containing metadata of the
	frames in the dataset, serialized List[types.FrameAnnotation].
	sequence_annotations_file: A zipped json file containing metadata of the
	sequences in the dataset, serialized List[types.SequenceAnnotation].
	subset_lists_file: A json file containing the lists of frames corresponding
	corresponding to different subsets (e.g. train/val/test) of the dataset;
	format: {subset: (sequence_name, frame_id, file_path)}.
	subsets: Restrict frames/sequences only to the given list of subsets
	as defined in subset_lists_file (see above).
	limit_to: Limit the dataset to the first #limit_to frames (after other
	filters have been applied).
	limit_sequences_to: Limit the dataset to the first
	#limit_sequences_to sequences (after other sequence filters have been
	applied but before frame-based filters).
	pick_sequence: A list of sequence names to restrict the dataset to.
	exclude_sequence: A list of the names of the sequences to exclude.
	limit_category_to: Restrict the dataset to the given list of categories.
	dataset_root: The root folder of the dataset; all the paths in jsons are
	specified relative to this root (but not json paths themselves).
	load_images: Enable loading the frame RGB data.
	load_depths: Enable loading the frame depth maps.
	load_depth_masks: Enable loading the frame depth map masks denoting the
	depth values used for evaluation (the points consistent across views).
	load_masks: Enable loading frame foreground masks.
	load_point_clouds: Enable loading sequence-level point clouds.
	max_points: Cap on the number of loaded points in the point cloud;
	if reached, they are randomly sampled without replacement.
	mask_images: Whether to mask the images with the loaded foreground masks;
	0 value is used for background.
	mask_depths: Whether to mask the depth maps with the loaded foreground
	masks; 0 value is used for background.
	image_height: The height of the returned images, masks, and depth maps;
	aspect ratio is preserved during cropping/resizing.
	image_width: The width of the returned images, masks, and depth maps;
	aspect ratio is preserved during cropping/resizing.
	box_crop: Enable cropping of the image around the bounding box inferred
	from the foreground region of the loaded segmentation mask; masks
	and depth maps are cropped accordingly; cameras are corrected.
	box_crop_mask_thr: The threshold used to separate pixels into foreground
	and background based on the foreground_probability mask; if no value
	is greater than this threshold, the loader lowers it and repeats.
	box_crop_context: The amount of additional padding added to each
	dimension of the cropping bounding box, relative to box size.
	remove_empty_masks: Removes the frames with no active foreground pixels
	in the segmentation mask after thresholding (see box_crop_mask_thr).
	n_frames_per_sequence: If > 0, randomly samples #n_frames_per_sequence
	frames in each sequences uniformly without replacement if it has
	more frames than that; applied before other frame-level filters.
	seed: The seed of the random generator sampling #n_frames_per_sequence
	random frames per sequence.
	sort_frames: Enable frame annotations sorting to group frames from the
	same sequences together and order them by timestamps
	eval_batches: A list of batches that form the evaluation set;
	list of batch-sized lists of indices corresponding to __getitem__
	of this class, thus it can be used directly as a batch sampler.
	eval_batch_index:
	( Optional[List[List[Union[Tuple[str, int, str], Tuple[str, int]]]] )
	A list of batches of frames described as (sequence_name, frame_idx)
	that can form the evaluation set, `eval_batches` will be set from this.

	"""

	frame_annotations_type: ClassVar[
	Type[types.FrameAnnotation]
	] = types.FrameAnnotation

	path_manager: Any = None
	frame_annotations_file: str = ""
	sequence_annotations_file: str = ""
	subset_lists_file: str = ""
	subsets: Optional[List[str]] = None
	limit_to: int = 0
	limit_sequences_to: int = 0
	pick_sequence: Tuple[str, ...] = ()
	exclude_sequence: Tuple[str, ...] = ()
	limit_category_to: Tuple[int, ...] = ()
	dataset_root: str = ""
	load_images: bool = True
	load_depths: bool = True
	load_depth_masks: bool = True
	load_masks: bool = True
	load_point_clouds: bool = False
	max_points: int = 0
	mask_images: bool = False
	mask_depths: bool = False
	image_height: Optional[int] = 800
	image_width: Optional[int] = 800
	box_crop: bool = True
	box_crop_mask_thr: float = 0.4
	box_crop_context: float = 0.3
	remove_empty_masks: bool = True
	n_frames_per_sequence: int = -1
	seed: int = 0
	sort_frames: bool = False
	eval_batches: Any = None
	eval_batch_index: Any = None
	# frame_annots: List[FrameAnnotsEntry] = field(init=False)
	# seq_annots: Dict[str, types.SequenceAnnotation] = field(init=False)

	def __post_init__(self) -> None:
	# pyre-fixme[16]: `JsonIndexDataset` has no attribute `subset_to_image_path`.
	self.subset_to_image_path = None
	self._load_frames()
	self._load_sequences()
	if self.sort_frames:
	self._sort_frames()
	self._load_subset_lists()
	self._filter_db() # also computes sequence indices
	self._extract_and_set_eval_batches()
	logger.info(str(self))

	def _extract_and_set_eval_batches(self):
	"""
	Sets eval_batches based on input eval_batch_index.
	"""
	if self.eval_batch_index is not None:
	if self.eval_batches is not None:
	raise ValueError(
	"Cannot define both eval_batch_index and eval_batches."
	)
	self.eval_batches = self.seq_frame_index_to_dataset_index(
	self.eval_batch_index
	)

	def join(self, other_datasets: Iterable[DatasetBase]) -> None:
	"""
	Join the dataset with other JsonIndexDataset objects.

	Args:
	other_datasets: A list of JsonIndexDataset objects to be joined
	into the current dataset.
	"""
	if not all(isinstance(d, JsonIndexDataset) for d in other_datasets):
	raise ValueError("This function can only join a list of JsonIndexDataset")
	# pyre-ignore[16]
	self.frame_annots.extend([fa for d in other_datasets for fa in d.frame_annots])
	# pyre-ignore[16]
	self.seq_annots.update(
	# https://gist.github.com/treyhunner/f35292e676efa0be1728
	functools.reduce(
	lambda a, b: {a, b},
	[d.seq_annots for d in other_datasets], # pyre-ignore[16]
	)
	)
	all_eval_batches = [
	self.eval_batches,
	# pyre-ignore
	*[d.eval_batches for d in other_datasets],
	]
	if not (
	all(ba is None for ba in all_eval_batches)
	or all(ba is not None for ba in all_eval_batches)
	):
	raise ValueError(
	"When joining datasets, either all joined datasets have to have their"
	" eval_batches defined, or all should have their eval batches undefined."
	)
	if self.eval_batches is not None:
	self.eval_batches = sum(all_eval_batches, [])
	self._invalidate_indexes(filter_seq_annots=True)

	def is_filtered(self) -> bool:
	"""
	Returns `True` in case the dataset has been filtered and thus some frame annotations
	stored on the disk might be missing in the dataset object.

	Returns:
	is_filtered: `True` if the dataset has been filtered, else `False`.
	"""
	return (
	self.remove_empty_masks
	or self.limit_to > 0
	or self.limit_sequences_to > 0
	or len(self.pick_sequence) > 0
	or len(self.exclude_sequence) > 0
	or len(self.limit_category_to) > 0
	or self.n_frames_per_sequence > 0
	)

	def seq_frame_index_to_dataset_index(
	self,
	seq_frame_index: List[List[Union[Tuple[str, int, str], Tuple[str, int]]]],
	allow_missing_indices: bool = False,
	remove_missing_indices: bool = False,
	suppress_missing_index_warning: bool = True,
	) -> List[List[Union[Optional[int], int]]]:
	"""
	Obtain indices into the dataset object given a list of frame ids.

	Args:
	seq_frame_index: The list of frame ids specified as
	`List[List[Tuple[sequence_name:str, frame_number:int]]]`. Optionally,
	Image paths relative to the dataset_root can be stored specified as well:
	`List[List[Tuple[sequence_name:str, frame_number:int, image_path:str]]]`
	allow_missing_indices: If `False`, throws an IndexError upon reaching the first
	entry from `seq_frame_index` which is missing in the dataset.
	Otherwise, depending on `remove_missing_indices`, either returns `None`
	in place of missing entries or removes the indices of missing entries.
	remove_missing_indices: Active when `allow_missing_indices=True`.
	If `False`, returns `None` in place of `seq_frame_index` entries that
	are not present in the dataset.
	If `True` removes missing indices from the returned indices.
	suppress_missing_index_warning:
	Active if `allow_missing_indices==True`. Suppressess a warning message
	in case an entry from `seq_frame_index` is missing in the dataset
	(expected in certain cases - e.g. when setting
	`self.remove_empty_masks=True`).

	Returns:
	dataset_idx: Indices of dataset entries corresponding to`seq_frame_index`.
	"""
	_dataset_seq_frame_n_index = {
	seq: {
	# pyre-ignore[16]
	self.frame_annots[idx]["frame_annotation"].frame_number: idx
	for idx in seq_idx
	}
	# pyre-ignore[16]
	for seq, seq_idx in self._seq_to_idx.items()
	}

	def _get_dataset_idx(
	seq_name: str, frame_no: int, path: Optional[str] = None
	) -> Optional[int]:
	idx_seq = _dataset_seq_frame_n_index.get(seq_name, None)
	idx = idx_seq.get(frame_no, None) if idx_seq is not None else None
	if idx is None:
	msg = (
	f"sequence_name={seq_name} / frame_number={frame_no}"
	" not in the dataset!"
	)
	if not allow_missing_indices:
	raise IndexError(msg)
	if not suppress_missing_index_warning:
	warnings.warn(msg)
	return idx
	if path is not None:
	# Check that the loaded frame path is consistent
	# with the one stored in self.frame_annots.
	assert os.path.normpath(
	# pyre-ignore[16]
	self.frame_annots[idx]["frame_annotation"].image.path
	) == os.path.normpath(
	path
	), f"Inconsistent frame indices {seq_name, frame_no, path}."
	return idx

	dataset_idx = [
	[_get_dataset_idx(*b) for b in batch] # pyre-ignore [6]
	for batch in seq_frame_index
	]

	if allow_missing_indices and remove_missing_indices:
	# remove all None indices, and also batches with only None entries
	valid_dataset_idx = [
	[b for b in batch if b is not None] for batch in dataset_idx
	]
	return [ # pyre-ignore[7]
	batch for batch in valid_dataset_idx if len(batch) > 0
	]

	return dataset_idx

	def subset_from_frame_index(
	self,
	frame_index: List[Union[Tuple[str, int], Tuple[str, int, str]]],
	allow_missing_indices: bool = True,
	) -> "JsonIndexDataset":
	"""
	Generate a dataset subset given the list of frames specified in `frame_index`.

	Args:
	frame_index: The list of frame indentifiers (as stored in the metadata)
	specified as `List[Tuple[sequence_name:str, frame_number:int]]`. Optionally,
	Image paths relative to the dataset_root can be stored specified as well:
	`List[Tuple[sequence_name:str, frame_number:int, image_path:str]]`,
	in the latter case, if imaga_path do not match the stored paths, an error
	is raised.
	allow_missing_indices: If `False`, throws an IndexError upon reaching the first
	entry from `frame_index` which is missing in the dataset.
	Otherwise, generates a subset consisting of frames entries that actually
	exist in the dataset.
	"""
	# Get the indices into the frame annots.
	dataset_indices = self.seq_frame_index_to_dataset_index(
	[frame_index],
	allow_missing_indices=self.is_filtered() and allow_missing_indices,
	)[0]
	valid_dataset_indices = [i for i in dataset_indices if i is not None]

	# Deep copy the whole dataset except frame_annots, which are large so we
	# deep copy only the requested subset of frame_annots.
	memo = {id(self.frame_annots): None} # pyre-ignore[16]
	dataset_new = copy.deepcopy(self, memo)
	dataset_new.frame_annots = copy.deepcopy(
	[self.frame_annots[i] for i in valid_dataset_indices]
	)

	# This will kill all unneeded sequence annotations.
	dataset_new._invalidate_indexes(filter_seq_annots=True)

	# Finally annotate the frame annotations with the name of the subset
	# stored in meta.
	for frame_annot in dataset_new.frame_annots:
	frame_annotation = frame_annot["frame_annotation"]
	if frame_annotation.meta is not None:
	frame_annot["subset"] = frame_annotation.meta.get("frame_type", None)

	# A sanity check - this will crash in case some entries from frame_index are missing
	# in dataset_new.
	valid_frame_index = [
	fi for fi, di in zip(frame_index, dataset_indices) if di is not None
	]
	dataset_new.seq_frame_index_to_dataset_index(
	[valid_frame_index], allow_missing_indices=False
	)

	return dataset_new

	def __str__(self) -> str:
	# pyre-ignore[16]
	return f"JsonIndexDataset #frames={len(self.frame_annots)}"

	def __len__(self) -> int:
	# pyre-ignore[16]
	return len(self.frame_annots)

	def _get_frame_type(self, entry: FrameAnnotsEntry) -> Optional[str]:
	return entry["subset"]

	def get_all_train_cameras(self) -> CamerasBase:
	"""
	Returns the cameras corresponding to all the known frames.
	"""
	logger.info("Loading all train cameras.")
	cameras = []
	# pyre-ignore[16]
	for frame_idx, frame_annot in enumerate(tqdm(self.frame_annots)):
	frame_type = self._get_frame_type(frame_annot)
	if frame_type is None:
	raise ValueError("subsets not loaded")
	if is_known_frame_scalar(frame_type):
	cameras.append(self[frame_idx].camera)
	return join_cameras_as_batch(cameras)

	def __getitem__(self, index) -> FrameData:
	# pyre-ignore[16]
	if index >= len(self.frame_annots):
	raise IndexError(f"index {index} out of range {len(self.frame_annots)}")

	entry = self.frame_annots[index]["frame_annotation"]
	# pyre-ignore[16]
	point_cloud = self.seq_annots[entry.sequence_name].point_cloud
	frame_data = FrameData(
	frame_number=_safe_as_tensor(entry.frame_number, torch.long),
	frame_timestamp=_safe_as_tensor(entry.frame_timestamp, torch.float),
	sequence_name=entry.sequence_name,
	sequence_category=self.seq_annots[entry.sequence_name].category,
	camera_quality_score=_safe_as_tensor(
	self.seq_annots[entry.sequence_name].viewpoint_quality_score,
	torch.float,
	),
	point_cloud_quality_score=_safe_as_tensor(
	point_cloud.quality_score, torch.float
	)
	if point_cloud is not None
	else None,
	)

	# The rest of the fields are optional
	frame_data.frame_type = self._get_frame_type(self.frame_annots[index])

	(
	frame_data.fg_probability,
	frame_data.mask_path,
	frame_data.bbox_xywh,
	clamp_bbox_xyxy,
	frame_data.crop_bbox_xywh,
	) = self._load_crop_fg_probability(entry)

	scale = 1.0
	if self.load_images and entry.image is not None:
	# original image size
	frame_data.image_size_hw = _safe_as_tensor(entry.image.size, torch.long)

	(
	frame_data.image_rgb,
	frame_data.image_path,
	frame_data.mask_crop,
	scale,
	) = self._load_crop_images(
	entry, frame_data.fg_probability, clamp_bbox_xyxy
	)

	if self.load_depths and entry.depth is not None:
	(
	frame_data.depth_map,
	frame_data.depth_path,
	frame_data.depth_mask,
	) = self._load_mask_depth(entry, clamp_bbox_xyxy, frame_data.fg_probability)

	if entry.viewpoint is not None:
	frame_data.camera = self._get_pytorch3d_camera(
	entry,
	scale,
	clamp_bbox_xyxy,
	)

	if self.load_point_clouds and point_cloud is not None:
	pcl_path = self._fix_point_cloud_path(point_cloud.path)
	frame_data.sequence_point_cloud = _load_pointcloud(
	self._local_path(pcl_path), max_points=self.max_points
	)
	frame_data.sequence_point_cloud_path = pcl_path

	return frame_data

	def _fix_point_cloud_path(self, path: str) -> str:
	"""
	Fix up a point cloud path from the dataset.
	Some files in Co3Dv2 have an accidental absolute path stored.
	"""
	unwanted_prefix = (
	"/large_experiments/p3/replay/datasets/co3d/co3d45k_220512/export_v23/"
	)
	if path.startswith(unwanted_prefix):
	path = path[len(unwanted_prefix) :]
	return os.path.join(self.dataset_root, path)

	def _load_crop_fg_probability(
	self, entry: types.FrameAnnotation
	) -> Tuple[
	Optional[torch.Tensor],
	Optional[str],
	Optional[torch.Tensor],
	Optional[torch.Tensor],
	Optional[torch.Tensor],
	]:
	fg_probability = None
	full_path = None
	bbox_xywh = None
	clamp_bbox_xyxy = None
	crop_box_xywh = None

	if (self.load_masks or self.box_crop) and entry.mask is not None:
	full_path = os.path.join(self.dataset_root, entry.mask.path)
	mask = _load_mask(self._local_path(full_path))

	if mask.shape[-2:] != entry.image.size:
	raise ValueError(
	f"bad mask size: {mask.shape[-2:]} vs {entry.image.size}!"
	)

	bbox_xywh = torch.tensor(_get_bbox_from_mask(mask, self.box_crop_mask_thr))

	if self.box_crop:
	clamp_bbox_xyxy = _clamp_box_to_image_bounds_and_round(
	_get_clamp_bbox(
	bbox_xywh,
	image_path=entry.image.path,
	box_crop_context=self.box_crop_context,
	),
	image_size_hw=tuple(mask.shape[-2:]),
	)
	crop_box_xywh = _bbox_xyxy_to_xywh(clamp_bbox_xyxy)

	mask = _crop_around_box(mask, clamp_bbox_xyxy, full_path)

	fg_probability, _, _ = self._resize_image(mask, mode="nearest")

	return fg_probability, full_path, bbox_xywh, clamp_bbox_xyxy, crop_box_xywh

	def _load_crop_images(
	self,
	entry: types.FrameAnnotation,
	fg_probability: Optional[torch.Tensor],
	clamp_bbox_xyxy: Optional[torch.Tensor],
	) -> Tuple[torch.Tensor, str, torch.Tensor, float]:
	assert self.dataset_root is not None and entry.image is not None
	path = os.path.join(self.dataset_root, entry.image.path)
	image_rgb = _load_image(self._local_path(path))

	if image_rgb.shape[-2:] != entry.image.size:
	raise ValueError(
	f"bad image size: {image_rgb.shape[-2:]} vs {entry.image.size}!"
	)

	if self.box_crop:
	assert clamp_bbox_xyxy is not None
	image_rgb = _crop_around_box(image_rgb, clamp_bbox_xyxy, path)

	image_rgb, scale, mask_crop = self._resize_image(image_rgb)

	if self.mask_images:
	assert fg_probability is not None
	image_rgb *= fg_probability

	return image_rgb, path, mask_crop, scale

	def _load_mask_depth(
	self,
	entry: types.FrameAnnotation,
	clamp_bbox_xyxy: Optional[torch.Tensor],
	fg_probability: Optional[torch.Tensor],
	) -> Tuple[torch.Tensor, str, torch.Tensor]:
	entry_depth = entry.depth
	assert entry_depth is not None
	path = os.path.join(self.dataset_root, entry_depth.path)
	depth_map = _load_depth(self._local_path(path), entry_depth.scale_adjustment)

	if self.box_crop:
	assert clamp_bbox_xyxy is not None
	depth_bbox_xyxy = _rescale_bbox(
	clamp_bbox_xyxy, entry.image.size, depth_map.shape[-2:]
	)
	depth_map = _crop_around_box(depth_map, depth_bbox_xyxy, path)

	depth_map, _, _ = self._resize_image(depth_map, mode="nearest")

	if self.mask_depths:
	assert fg_probability is not None
	depth_map *= fg_probability

	if self.load_depth_masks:
	assert entry_depth.mask_path is not None
	mask_path = os.path.join(self.dataset_root, entry_depth.mask_path)
	depth_mask = _load_depth_mask(self._local_path(mask_path))

	if self.box_crop:
	assert clamp_bbox_xyxy is not None
	depth_mask_bbox_xyxy = _rescale_bbox(
	clamp_bbox_xyxy, entry.image.size, depth_mask.shape[-2:]
	)
	depth_mask = _crop_around_box(
	depth_mask, depth_mask_bbox_xyxy, mask_path
	)

	depth_mask, _, _ = self._resize_image(depth_mask, mode="nearest")
	else:
	depth_mask = torch.ones_like(depth_map)

	return depth_map, path, depth_mask

	def _get_pytorch3d_camera(
	self,
	entry: types.FrameAnnotation,
	scale: float,
	clamp_bbox_xyxy: Optional[torch.Tensor],
	) -> PerspectiveCameras:
	entry_viewpoint = entry.viewpoint
	assert entry_viewpoint is not None
	# principal point and focal length
	principal_point = torch.tensor(
	entry_viewpoint.principal_point, dtype=torch.float
	)
	focal_length = torch.tensor(entry_viewpoint.focal_length, dtype=torch.float)

	half_image_size_wh_orig = (
	torch.tensor(list(reversed(entry.image.size)), dtype=torch.float) / 2.0
	)

	# first, we convert from the dataset's NDC convention to pixels
	format = entry_viewpoint.intrinsics_format
	if format.lower() == "ndc_norm_image_bounds":
	# this is e.g. currently used in CO3D for storing intrinsics
	rescale = half_image_size_wh_orig
	elif format.lower() == "ndc_isotropic":
	rescale = half_image_size_wh_orig.min()
	else:
	raise ValueError(f"Unknown intrinsics format: {format}")

	# principal point and focal length in pixels
	principal_point_px = half_image_size_wh_orig - principal_point * rescale
	focal_length_px = focal_length * rescale
	if self.box_crop:
	assert clamp_bbox_xyxy is not None
	principal_point_px -= clamp_bbox_xyxy[:2]

	# now, convert from pixels to PyTorch3D v0.5+ NDC convention
	if self.image_height is None or self.image_width is None:
	out_size = list(reversed(entry.image.size))
	else:
	out_size = [self.image_width, self.image_height]

	half_image_size_output = torch.tensor(out_size, dtype=torch.float) / 2.0
	half_min_image_size_output = half_image_size_output.min()

	# rescaled principal point and focal length in ndc
	principal_point = (
	half_image_size_output - principal_point_px * scale
	) / half_min_image_size_output
	focal_length = focal_length_px * scale / half_min_image_size_output

	return PerspectiveCameras(
	focal_length=focal_length[None],
	principal_point=principal_point[None],
	R=torch.tensor(entry_viewpoint.R, dtype=torch.float)[None],
	T=torch.tensor(entry_viewpoint.T, dtype=torch.float)[None],
	)

	def _load_frames(self) -> None:
	logger.info(f"Loading Co3D frames from {self.frame_annotations_file}.")
	local_file = self._local_path(self.frame_annotations_file)
	with gzip.open(local_file, "rt", encoding="utf8") as zipfile:
	frame_annots_list = types.load_dataclass(
	zipfile, List[self.frame_annotations_type]
	)
	if not frame_annots_list:
	raise ValueError("Empty dataset!")
	# pyre-ignore[16]
	self.frame_annots = [
	FrameAnnotsEntry(frame_annotation=a, subset=None) for a in frame_annots_list
	]

	def _load_sequences(self) -> None:
	logger.info(f"Loading Co3D sequences from {self.sequence_annotations_file}.")
	local_file = self._local_path(self.sequence_annotations_file)
	with gzip.open(local_file, "rt", encoding="utf8") as zipfile:
	seq_annots = types.load_dataclass(zipfile, List[types.SequenceAnnotation])
	if not seq_annots:
	raise ValueError("Empty sequences file!")
	# pyre-ignore[16]
	self.seq_annots = {entry.sequence_name: entry for entry in seq_annots}

	def _load_subset_lists(self) -> None:
	logger.info(f"Loading Co3D subset lists from {self.subset_lists_file}.")
	if not self.subset_lists_file:
	return

	with open(self._local_path(self.subset_lists_file), "r") as f:
	subset_to_seq_frame = json.load(f)

	frame_path_to_subset = {
	path: subset
	for subset, frames in subset_to_seq_frame.items()
	for _, _, path in frames
	}
	# pyre-ignore[16]
	for frame in self.frame_annots:
	frame["subset"] = frame_path_to_subset.get(
	frame["frame_annotation"].image.path, None
	)
	if frame["subset"] is None:
	warnings.warn(
	"Subset lists are given but don't include "
	+ frame["frame_annotation"].image.path
	)

	def _sort_frames(self) -> None:
	# Sort frames to have them grouped by sequence, ordered by timestamp
	# pyre-ignore[16]
	self.frame_annots = sorted(
	self.frame_annots,
	key=lambda f: (
	f["frame_annotation"].sequence_name,
	f["frame_annotation"].frame_timestamp or 0,
	),
	)

	def _filter_db(self) -> None:
	if self.remove_empty_masks:
	logger.info("Removing images with empty masks.")
	# pyre-ignore[16]
	old_len = len(self.frame_annots)

	msg = "remove_empty_masks needs every MaskAnnotation.mass to be set."

	def positive_mass(frame_annot: types.FrameAnnotation) -> bool:
	mask = frame_annot.mask
	if mask is None:
	return False
	if mask.mass is None:
	raise ValueError(msg)
	return mask.mass > 1

	self.frame_annots = [
	frame
	for frame in self.frame_annots
	if positive_mass(frame["frame_annotation"])
	]
	logger.info("... filtered %d -> %d" % (old_len, len(self.frame_annots)))

	# this has to be called after joining with categories!!
	subsets = self.subsets
	if subsets:
	if not self.subset_lists_file:
	raise ValueError(
	"Subset filter is on but subset_lists_file was not given"
	)

	logger.info(f"Limiting Co3D dataset to the '{subsets}' subsets.")

	# truncate the list of subsets to the valid one
	self.frame_annots = [
	entry for entry in self.frame_annots if entry["subset"] in subsets
	]
	if len(self.frame_annots) == 0:
	raise ValueError(f"There are no frames in the '{subsets}' subsets!")

	self._invalidate_indexes(filter_seq_annots=True)

	if len(self.limit_category_to) > 0:
	logger.info(f"Limiting dataset to categories: {self.limit_category_to}")
	# pyre-ignore[16]
	self.seq_annots = {
	name: entry
	for name, entry in self.seq_annots.items()
	if entry.category in self.limit_category_to
	}

	# sequence filters
	for prefix in ("pick", "exclude"):
	orig_len = len(self.seq_annots)
	attr = f"{prefix}_sequence"
	arr = getattr(self, attr)
	if len(arr) > 0:
	logger.info(f"{attr}: {str(arr)}")
	self.seq_annots = {
	name: entry
	for name, entry in self.seq_annots.items()
	if (name in arr) == (prefix == "pick")
	}
	logger.info("... filtered %d -> %d" % (orig_len, len(self.seq_annots)))

	if self.limit_sequences_to > 0:
	self.seq_annots = dict(
	islice(self.seq_annots.items(), self.limit_sequences_to)
	)

	# retain only frames from retained sequences
	self.frame_annots = [
	f
	for f in self.frame_annots
	if f["frame_annotation"].sequence_name in self.seq_annots
	]

	self._invalidate_indexes()

	if self.n_frames_per_sequence > 0:
	logger.info(f"Taking max {self.n_frames_per_sequence} per sequence.")
	keep_idx = []
	# pyre-ignore[16]
	for seq, seq_indices in self._seq_to_idx.items():
	# infer the seed from the sequence name, this is reproducible
	# and makes the selection differ for different sequences
	seed = _seq_name_to_seed(seq) + self.seed
	seq_idx_shuffled = random.Random(seed).sample(
	sorted(seq_indices), len(seq_indices)
	)
	keep_idx.extend(seq_idx_shuffled[: self.n_frames_per_sequence])

	logger.info(
	"... filtered %d -> %d" % (len(self.frame_annots), len(keep_idx))
	)
	self.frame_annots = [self.frame_annots[i] for i in keep_idx]
	self._invalidate_indexes(filter_seq_annots=False)
	# sequences are not decimated, so self.seq_annots is valid

	if self.limit_to > 0 and self.limit_to < len(self.frame_annots):
	logger.info(
	"limit_to: filtered %d -> %d" % (len(self.frame_annots), self.limit_to)
	)
	self.frame_annots = self.frame_annots[: self.limit_to]
	self._invalidate_indexes(filter_seq_annots=True)

	def _invalidate_indexes(self, filter_seq_annots: bool = False) -> None:
	# update _seq_to_idx and filter seq_meta according to frame_annots change
	# if filter_seq_annots, also uldates seq_annots based on the changed _seq_to_idx
	self._invalidate_seq_to_idx()

	if filter_seq_annots:
	# pyre-ignore[16]
	self.seq_annots = {
	k: v
	for k, v in self.seq_annots.items()
	# pyre-ignore[16]
	if k in self._seq_to_idx
	}

	def _invalidate_seq_to_idx(self) -> None:
	seq_to_idx = defaultdict(list)
	# pyre-ignore[16]
	for idx, entry in enumerate(self.frame_annots):
	seq_to_idx[entry["frame_annotation"].sequence_name].append(idx)
	# pyre-ignore[16]
	self._seq_to_idx = seq_to_idx

	def _resize_image(
	self, image, mode="bilinear"
	) -> Tuple[torch.Tensor, float, torch.Tensor]:
	image_height, image_width = self.image_height, self.image_width
	if image_height is None or image_width is None:
	# skip the resizing
	imre_ = torch.from_numpy(image)
	return imre_, 1.0, torch.ones_like(imre_[:1])
	# takes numpy array, returns pytorch tensor
	minscale = min(
	image_height / image.shape[-2],
	image_width / image.shape[-1],
	)
	imre = torch.nn.functional.interpolate(
	torch.from_numpy(image)[None],
	scale_factor=minscale,
	mode=mode,
	align_corners=False if mode == "bilinear" else None,
	recompute_scale_factor=True,
	)[0]
	# pyre-fixme[19]: Expected 1 positional argument.
	imre_ = torch.zeros(image.shape[0], self.image_height, self.image_width)
	imre_[:, 0 : imre.shape[1], 0 : imre.shape[2]] = imre
	# pyre-fixme[6]: For 2nd param expected `int` but got `Optional[int]`.
	# pyre-fixme[6]: For 3rd param expected `int` but got `Optional[int]`.
	mask = torch.zeros(1, self.image_height, self.image_width)
	mask[:, 0 : imre.shape[1], 0 : imre.shape[2]] = 1.0
	return imre_, minscale, mask

	def _local_path(self, path: str) -> str:
	if self.path_manager is None:
	return path
	return self.path_manager.get_local_path(path)

	def get_frame_numbers_and_timestamps(
	self, idxs: Sequence[int]
	) -> List[Tuple[int, float]]:
	out: List[Tuple[int, float]] = []
	for idx in idxs:
	# pyre-ignore[16]
	frame_annotation = self.frame_annots[idx]["frame_annotation"]
	out.append(
	(frame_annotation.frame_number, frame_annotation.frame_timestamp)
	)
	return out

	def category_to_sequence_names(self) -> Dict[str, List[str]]:
	c2seq = defaultdict(list)
	# pyre-ignore
	for sequence_name, sa in self.seq_annots.items():
	c2seq[sa.category].append(sequence_name)
	return dict(c2seq)

	def get_eval_batches(self) -> Optional[List[List[int]]]:
	return self.eval_batches


	def _seq_name_to_seed(seq_name) -> int:
	return int(hashlib.sha1(seq_name.encode("utf-8")).hexdigest(), 16)


	def _load_image(path) -> np.ndarray:
	with Image.open(path) as pil_im:
	im = np.array(pil_im.convert("RGB"))
	im = im.transpose((2, 0, 1))
	im = im.astype(np.float32) / 255.0
	return im


	def _load_16big_png_depth(depth_png) -> np.ndarray:
	with Image.open(depth_png) as depth_pil:
	# the image is stored with 16-bit depth but PIL reads it as I (32 bit).
	# we cast it to uint16, then reinterpret as float16, then cast to float32
	depth = (
	np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
	.astype(np.float32)
	.reshape((depth_pil.size[1], depth_pil.size[0]))
	)
	return depth


	def _load_1bit_png_mask(file: str) -> np.ndarray:
	with Image.open(file) as pil_im:
	mask = (np.array(pil_im.convert("L")) > 0.0).astype(np.float32)
	return mask


	def _load_depth_mask(path: str) -> np.ndarray:
	if not path.lower().endswith(".png"):
	raise ValueError('unsupported depth mask file name "%s"' % path)
	m = _load_1bit_png_mask(path)
	return m[None] # fake feature channel


	def _load_depth(path, scale_adjustment) -> np.ndarray:
	if not path.lower().endswith(".png"):
	raise ValueError('unsupported depth file name "%s"' % path)

	d = _load_16big_png_depth(path) * scale_adjustment
	d[~np.isfinite(d)] = 0.0
	return d[None] # fake feature channel


	def _load_mask(path) -> np.ndarray:
	with Image.open(path) as pil_im:
	mask = np.array(pil_im)
	mask = mask.astype(np.float32) / 255.0
	return mask[None] # fake feature channel


	def _get_1d_bounds(arr) -> Tuple[int, int]:
	nz = np.flatnonzero(arr)
	return nz[0], nz[-1] + 1


	def _get_bbox_from_mask(
	mask, thr, decrease_quant: float = 0.05
	) -> Tuple[int, int, int, int]:
	# bbox in xywh
	masks_for_box = np.zeros_like(mask)
	while masks_for_box.sum() <= 1.0:
	masks_for_box = (mask > thr).astype(np.float32)
	thr -= decrease_quant
	if thr <= 0.0:
	warnings.warn(f"Empty masks_for_bbox (thr={thr}) => using full image.")

	x0, x1 = _get_1d_bounds(masks_for_box.sum(axis=-2))
	y0, y1 = _get_1d_bounds(masks_for_box.sum(axis=-1))

	return x0, y0, x1 - x0, y1 - y0


	def _get_clamp_bbox(
	bbox: torch.Tensor,
	box_crop_context: float = 0.0,
	image_path: str = "",
	) -> torch.Tensor:
	# box_crop_context: rate of expansion for bbox
	# returns possibly expanded bbox xyxy as float

	bbox = bbox.clone() # do not edit bbox in place

	# increase box size
	if box_crop_context > 0.0:
	c = box_crop_context
	bbox = bbox.float()
	bbox[0] -= bbox[2] * c / 2
	bbox[1] -= bbox[3] * c / 2
	bbox[2] += bbox[2] * c
	bbox[3] += bbox[3] * c

	if (bbox[2:] <= 1.0).any():
	raise ValueError(
	f"squashed image {image_path}!! The bounding box contains no pixels."
	)

	bbox[2:] = torch.clamp(bbox[2:], 2) # set min height, width to 2 along both axes
	bbox_xyxy = _bbox_xywh_to_xyxy(bbox, clamp_size=2)

	return bbox_xyxy


	def _crop_around_box(tensor, bbox, impath: str = ""):
	# bbox is xyxy, where the upper bound is corrected with +1
	bbox = _clamp_box_to_image_bounds_and_round(
	bbox,
	image_size_hw=tensor.shape[-2:],
	)
	tensor = tensor[..., bbox[1] : bbox[3], bbox[0] : bbox[2]]
	assert all(c > 0 for c in tensor.shape), f"squashed image {impath}"
	return tensor


	def _clamp_box_to_image_bounds_and_round(
	bbox_xyxy: torch.Tensor,
	image_size_hw: Tuple[int, int],
	) -> torch.LongTensor:
	bbox_xyxy = bbox_xyxy.clone()
	bbox_xyxy[[0, 2]] = torch.clamp(bbox_xyxy[[0, 2]], 0, image_size_hw[-1])
	bbox_xyxy[[1, 3]] = torch.clamp(bbox_xyxy[[1, 3]], 0, image_size_hw[-2])
	if not isinstance(bbox_xyxy, torch.LongTensor):
	bbox_xyxy = bbox_xyxy.round().long()
	return bbox_xyxy # pyre-ignore [7]


	def _rescale_bbox(bbox: torch.Tensor, orig_res, new_res) -> torch.Tensor:
	assert bbox is not None
	assert np.prod(orig_res) > 1e-8
	# average ratio of dimensions
	rel_size = (new_res[0] / orig_res[0] + new_res[1] / orig_res[1]) / 2.0
	return bbox * rel_size


	def _bbox_xyxy_to_xywh(xyxy: torch.Tensor) -> torch.Tensor:
	wh = xyxy[2:] - xyxy[:2]
	xywh = torch.cat([xyxy[:2], wh])
	return xywh


	def _bbox_xywh_to_xyxy(
	xywh: torch.Tensor, clamp_size: Optional[int] = None
	) -> torch.Tensor:
	xyxy = xywh.clone()
	if clamp_size is not None:
	xyxy[2:] = torch.clamp(xyxy[2:], clamp_size)
	xyxy[2:] += xyxy[:2]
	return xyxy


	def _safe_as_tensor(data, dtype):
	if data is None:
	return None
	return torch.tensor(data, dtype=dtype)


	# NOTE this cache is per-worker; they are implemented as processes.
	# each batch is loaded and collated by a single worker;
	# since sequences tend to co-occur within batches, this is useful.
	@functools.lru_cache(maxsize=256)
	def _load_pointcloud(pcl_path: Union[str, Path], max_points: int = 0) -> Pointclouds:
	pcl = IO().load_pointcloud(pcl_path)
	if max_points > 0:
	pcl = pcl.subsample(max_points)

	return pcl