import string import h5py import torch from ..datasets.base_dataset import collate from ..settings import DATA_PATH from ..utils.tensor import batch_to_device from .base_model import BaseModel from .utils.misc import pad_to_length def pad_local_features(pred: dict, seq_l: int): pred["keypoints"] = pad_to_length( pred["keypoints"], seq_l, -2, mode="random_c", ) if "keypoint_scores" in pred.keys(): pred["keypoint_scores"] = pad_to_length( pred["keypoint_scores"], seq_l, -1, mode="zeros" ) if "descriptors" in pred.keys(): pred["descriptors"] = pad_to_length( pred["descriptors"], seq_l, -2, mode="random" ) if "scales" in pred.keys(): pred["scales"] = pad_to_length(pred["scales"], seq_l, -1, mode="zeros") if "oris" in pred.keys(): pred["oris"] = pad_to_length(pred["oris"], seq_l, -1, mode="zeros") if "depth_keypoints" in pred.keys(): pred["depth_keypoints"] = pad_to_length( pred["depth_keypoints"], seq_l, -1, mode="zeros" ) if "valid_depth_keypoints" in pred.keys(): pred["valid_depth_keypoints"] = pad_to_length( pred["valid_depth_keypoints"], seq_l, -1, mode="zeros" ) return pred def pad_line_features(pred, seq_l: int = None): raise NotImplementedError def recursive_load(grp, pkeys): return { k: torch.from_numpy(grp[k].__array__()) if isinstance(grp[k], h5py.Dataset) else recursive_load(grp[k], list(grp.keys())) for k in pkeys } class CacheLoader(BaseModel): default_conf = { "path": "???", # can be a format string like exports/{scene}/ "data_keys": None, # load all keys "device": None, # load to same device as data "trainable": False, "add_data_path": True, "collate": True, "scale": ["keypoints", "lines", "orig_lines"], "padding_fn": None, "padding_length": None, # required for batching! "numeric_type": "float32", # [None, "float16", "float32", "float64"] } required_data_keys = ["name"] # we need an identifier def _init(self, conf): self.hfiles = {} self.padding_fn = conf.padding_fn if self.padding_fn is not None: self.padding_fn = eval(self.padding_fn) self.numeric_dtype = { None: None, "float16": torch.float16, "float32": torch.float32, "float64": torch.float64, }[conf.numeric_type] def _forward(self, data): preds = [] device = self.conf.device if not device: devices = set( [v.device for v in data.values() if isinstance(v, torch.Tensor)] ) if len(devices) == 0: device = "cpu" else: assert len(devices) == 1 device = devices.pop() var_names = [x[1] for x in string.Formatter().parse(self.conf.path) if x[1]] for i, name in enumerate(data["name"]): fpath = self.conf.path.format(**{k: data[k][i] for k in var_names}) if self.conf.add_data_path: fpath = DATA_PATH / fpath hfile = h5py.File(str(fpath), "r") grp = hfile[name] pkeys = ( self.conf.data_keys if self.conf.data_keys is not None else grp.keys() ) pred = recursive_load(grp, pkeys) if self.numeric_dtype is not None: pred = { k: v if not isinstance(v, torch.Tensor) or not torch.is_floating_point(v) else v.to(dtype=self.numeric_dtype) for k, v in pred.items() } pred = batch_to_device(pred, device) for k, v in pred.items(): for pattern in self.conf.scale: if k.startswith(pattern): view_idx = k.replace(pattern, "") scales = ( data["scales"] if len(view_idx) == 0 else data[f"view{view_idx}"]["scales"] ) pred[k] = pred[k] * scales[i] # use this function to fix number of keypoints etc. if self.padding_fn is not None: pred = self.padding_fn(pred, self.conf.padding_length) preds.append(pred) hfile.close() if self.conf.collate: return batch_to_device(collate(preds), device) else: assert len(preds) == 1 return batch_to_device(preds[0], device) def loss(self, pred, data): raise NotImplementedError