Spaces:
Running
Running
File size: 1,804 Bytes
404d2af 8b973ee 404d2af 8b973ee 404d2af 8b973ee |
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 |
import os
import numpy as np
import h5py
import cv2
class standard_reader:
def __init__(self, config):
self.raw_dir = config["rawdata_dir"]
self.dataset = h5py.File(config["dataset_dir"], "r")
self.num_kpt = config["num_kpt"]
def run(self, index):
K1, K2 = np.asarray(self.dataset["K1"][str(index)]), np.asarray(
self.dataset["K2"][str(index)]
)
R = np.asarray(self.dataset["R"][str(index)])
t = np.asarray(self.dataset["T"][str(index)])
t = t / np.sqrt((t**2).sum())
desc1, desc2 = (
self.dataset["desc1"][str(index)][()][: self.num_kpt],
self.dataset["desc2"][str(index)][()][: self.num_kpt],
)
x1, x2 = (
self.dataset["kpt1"][str(index)][()][: self.num_kpt],
self.dataset["kpt2"][str(index)][()][: self.num_kpt],
)
e, f = self.dataset["e"][str(index)][()], self.dataset["f"][str(index)][()]
img1_path, img2_path = (
self.dataset["img_path1"][str(index)][()][0].decode(),
self.dataset["img_path2"][str(index)][()][0].decode(),
)
img1, img2 = cv2.imread(os.path.join(self.raw_dir, img1_path)), cv2.imread(
os.path.join(self.raw_dir, img2_path)
)
info = {
"index": index,
"K1": K1,
"K2": K2,
"R": R,
"t": t,
"x1": x1,
"x2": x2,
"desc1": desc1,
"desc2": desc2,
"img1": img1,
"img2": img2,
"e": e,
"f": f,
"r_gt": R,
"t_gt": t,
}
return info
def close(self):
self.dataset.close()
def __len__(self):
return len(self.dataset["K1"])
|