Spaces:
Running
Running
File size: 2,661 Bytes
4d4dd90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
"""
Export the predictions of a model for a given dataloader (e.g. ImageFolder).
Use a standalone script with `python3 -m dsfm.scipts.export_predictions dir`
or call from another script.
"""
from pathlib import Path
import h5py
import numpy as np
import torch
from tqdm import tqdm
from .tensor import batch_to_device
@torch.no_grad()
def export_predictions(
loader,
model,
output_file,
as_half=False,
keys="*",
callback_fn=None,
optional_keys=[],
):
assert keys == "*" or isinstance(keys, (tuple, list))
Path(output_file).parent.mkdir(exist_ok=True, parents=True)
hfile = h5py.File(str(output_file), "w")
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device).eval()
for data_ in tqdm(loader):
data = batch_to_device(data_, device, non_blocking=True)
pred = model(data)
if callback_fn is not None:
pred = {**callback_fn(pred, data), **pred}
if keys != "*":
if len(set(keys) - set(pred.keys())) > 0:
raise ValueError(f"Missing key {set(keys) - set(pred.keys())}")
pred = {k: v for k, v in pred.items() if k in keys + optional_keys}
assert len(pred) > 0
# renormalization
for k in pred.keys():
if k.startswith("keypoints"):
idx = k.replace("keypoints", "")
scales = 1.0 / (
data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
)
pred[k] = pred[k] * scales[None]
if k.startswith("lines"):
idx = k.replace("lines", "")
scales = 1.0 / (
data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
)
pred[k] = pred[k] * scales[None]
if k.startswith("orig_lines"):
idx = k.replace("orig_lines", "")
scales = 1.0 / (
data["scales"] if len(idx) == 0 else data[f"view{idx}"]["scales"]
)
pred[k] = pred[k] * scales[None]
pred = {k: v[0].cpu().numpy() for k, v in pred.items()}
if as_half:
for k in pred:
dt = pred[k].dtype
if (dt == np.float32) and (dt != np.float16):
pred[k] = pred[k].astype(np.float16)
try:
name = data["name"][0]
grp = hfile.create_group(name)
for k, v in pred.items():
grp.create_dataset(k, data=v)
except RuntimeError:
continue
del pred
hfile.close()
return output_file
|