Vincentqyw
fix: roma
c74a070
raw
history blame
5.86 kB
import argparse
import glob
import math
import subprocess
import numpy as np
import os
import tqdm
import torch
import torch.nn as nn
import cv2
from darkfeat import DarkFeat
from utils import matching
def darkfeat_pre(img, cuda):
H, W = img.shape[0], img.shape[1]
inp = img.copy()
inp = inp.transpose(2, 0, 1)
inp = torch.from_numpy(inp)
inp = torch.autograd.Variable(inp).view(1, 3, H, W)
if cuda:
inp = inp.cuda()
return inp
if __name__ == "__main__":
# Parse command line arguments.
parser = argparse.ArgumentParser()
parser.add_argument("--H", type=int, default=int(640))
parser.add_argument("--W", type=int, default=int(960))
parser.add_argument("--histeq", action="store_true")
parser.add_argument("--model_path", type=str)
parser.add_argument("--dataset_dir", type=str, default="/data/hyz/MID/")
opt = parser.parse_args()
sizer = (opt.W, opt.H)
focallength_x = 4.504986436499113e03 / (6744 / sizer[0])
focallength_y = 4.513311442889859e03 / (4502 / sizer[1])
K = np.eye(3)
K[0, 0] = focallength_x
K[1, 1] = focallength_y
K[0, 2] = 3.363322177533149e03 / (6744 / sizer[0]) # * 0.5
K[1, 2] = 2.291824660547715e03 / (4502 / sizer[1]) # * 0.5
Kinv = np.linalg.inv(K)
Kinvt = np.transpose(Kinv)
cuda = True
if cuda:
darkfeat = DarkFeat(opt.model_path).cuda().eval()
for scene in ["Indoor", "Outdoor"]:
base_save = "./result/" + scene + "/"
dir_base = opt.dataset_dir + "/" + scene + "/"
pair_list = sorted(os.listdir(dir_base))
for pair in tqdm.tqdm(pair_list):
opention = 1
if scene == "Outdoor":
pass
else:
if int(pair[4::]) <= 17:
opention = 0
else:
pass
name = []
files = sorted(os.listdir(dir_base + pair))
for file_ in files:
if file_.endswith(".cr2"):
name.append(file_[0:9])
ISO = [
"00100",
"00200",
"00400",
"00800",
"01600",
"03200",
"06400",
"12800",
]
if opention == 1:
Shutter_speed = ["0.005", "0.01", "0.025", "0.05", "0.17", "0.5"]
else:
Shutter_speed = ["0.01", "0.02", "0.05", "0.1", "0.3", "1"]
E_GT = np.load(dir_base + pair + "/GT_Correspondence/" + "E_estimated.npy")
F_GT = np.dot(np.dot(Kinvt, E_GT), Kinv)
R_GT = np.load(dir_base + pair + "/GT_Correspondence/" + "R_GT.npy")
t_GT = np.load(dir_base + pair + "/GT_Correspondence/" + "T_GT.npy")
id0, id1 = sorted(
[int(i.split("/")[-1]) for i in glob.glob(f"{dir_base+pair}/?????")]
)
cnt = 0
for iso in ISO:
for ex in Shutter_speed:
dark_name1 = name[0] + iso + "_" + ex + "_" + scene + ".npy"
dark_name2 = name[1] + iso + "_" + ex + "_" + scene + ".npy"
if not opt.histeq:
dst_T1_None = (
f"{dir_base}{pair}/{id0:05d}-npy-nohisteq/{dark_name1}"
)
dst_T2_None = (
f"{dir_base}{pair}/{id1:05d}-npy-nohisteq/{dark_name2}"
)
img1_orig_None = np.load(dst_T1_None)
img2_orig_None = np.load(dst_T2_None)
dir_save = base_save + pair + "/None/"
img_input1 = darkfeat_pre(
img1_orig_None.astype("float32") / 255.0, cuda
)
img_input2 = darkfeat_pre(
img2_orig_None.astype("float32") / 255.0, cuda
)
else:
dst_T1_histeq = f"{dir_base}{pair}/{id0:05d}-npy/{dark_name1}"
dst_T2_histeq = f"{dir_base}{pair}/{id1:05d}-npy/{dark_name2}"
img1_orig_histeq = np.load(dst_T1_histeq)
img2_orig_histeq = np.load(dst_T2_histeq)
dir_save = base_save + pair + "/HistEQ/"
img_input1 = darkfeat_pre(
img1_orig_histeq.astype("float32") / 255.0, cuda
)
img_input2 = darkfeat_pre(
img2_orig_histeq.astype("float32") / 255.0, cuda
)
result1 = darkfeat({"image": img_input1})
result2 = darkfeat({"image": img_input2})
mkpts0, mkpts1, _ = matching.match_descriptors(
cv2.KeyPoint_convert(
result1["keypoints"].detach().cpu().float().numpy()
),
result1["descriptors"].detach().cpu().numpy(),
cv2.KeyPoint_convert(
result2["keypoints"].detach().cpu().float().numpy()
),
result2["descriptors"].detach().cpu().numpy(),
ORB=False,
)
POINT_1_dir = dir_save + f"DarkFeat/POINT_1/"
POINT_2_dir = dir_save + f"DarkFeat/POINT_2/"
subprocess.check_output(["mkdir", "-p", POINT_1_dir])
subprocess.check_output(["mkdir", "-p", POINT_2_dir])
np.save(POINT_1_dir + dark_name1[0:-3] + "npy", mkpts0)
np.save(POINT_2_dir + dark_name2[0:-3] + "npy", mkpts1)