mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
295 lines
14 KiB
Python
295 lines
14 KiB
Python
import cv2
|
|
from matplotlib import pyplot as plt
|
|
import numpy as np
|
|
import os
|
|
import json
|
|
import subprocess
|
|
import argparse
|
|
|
|
class ContourTooSmallException(Exception):
|
|
pass
|
|
|
|
class Image:
|
|
def __init__(self, image, threshold = 5):
|
|
self.image = image.copy()
|
|
self.threshold = threshold
|
|
self.gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
|
|
_, self.thresholded = cv2.threshold(self.gray, threshold, 255, cv2.THRESH_BINARY)
|
|
self.contours, _ = cv2.findContours(self.thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
|
|
|
class Contour:
|
|
def __init__(self, contour, well_image_threshold_index = None, min_area = 50.0):
|
|
self.contour = contour
|
|
self.area = cv2.contourArea(contour)
|
|
if self.area < min_area:
|
|
raise ContourTooSmallException(f"Contour area {self.area} is too small.")
|
|
self.centroid = np.mean(contour, axis=0)[0]
|
|
self.relative_position = None
|
|
self.match = None
|
|
self.score = float('inf')
|
|
self.well_image_threshold_index = well_image_threshold_index
|
|
|
|
def overlap(self, contour):
|
|
image_shape = (960, 960)
|
|
mask1 = np.zeros(image_shape, dtype=np.uint8)
|
|
mask2 = np.zeros(image_shape, dtype=np.uint8)
|
|
cv2.drawContours(mask1, [self.contour], -1, 255, thickness=cv2.FILLED)
|
|
cv2.drawContours(mask2, [contour.contour], -1, 255, thickness=cv2.FILLED)
|
|
overlap_mask = cv2.bitwise_and(mask1, mask2)
|
|
overlap_area = np.sum(overlap_mask) / 255
|
|
return overlap_area > 0
|
|
|
|
def get_matching_score(self, contour):
|
|
dist_shape = cv2.matchShapes(self.contour, contour.contour, cv2.CONTOURS_MATCH_I2, 0.0)
|
|
dist_spatial = np.linalg.norm(np.array(self.centroid) - np.array(contour.centroid))
|
|
return 0.7 * dist_shape + 0.3 * dist_spatial
|
|
|
|
def match_contour(self, contour, score):
|
|
self.match = contour
|
|
self.score = score
|
|
|
|
def unmatch_contour(self):
|
|
self.match = None
|
|
self.score = float('inf')
|
|
|
|
def get_points(self):
|
|
return self.contour[:, 0, :]
|
|
|
|
|
|
class AlignImages:
|
|
def __init__(self, well_image, fast_image):
|
|
self.fast_image = Image(fast_image)
|
|
self.fast_contours = []
|
|
for contour in self.fast_image.contours:
|
|
try:
|
|
contour_obj = Contour(contour)
|
|
self.fast_contours.append(contour_obj)
|
|
except ContourTooSmallException as e:
|
|
continue
|
|
self.matched_well_contours = [None] * len(self.fast_contours)
|
|
self.well_image = Image(well_image)
|
|
self.well_images = []
|
|
self.well_thresholds = (230, 235, 240, 245, 250)
|
|
for threshold in self.well_thresholds:
|
|
self.well_images.append(Image(well_image, threshold))
|
|
self.image_height, self.image_width = self.fast_image.image.shape[:2]
|
|
|
|
def all_contours_matched(self):
|
|
for contour in self.fast_contours:
|
|
if contour.match == None:
|
|
return False
|
|
return True
|
|
|
|
def existing_contour_match(self, contour):
|
|
for i, existing_contour in enumerate(self.matched_well_contours):
|
|
if existing_contour != None and contour.overlap(existing_contour):
|
|
return i, self.fast_contours[i]
|
|
return None, None
|
|
|
|
def delete_unmatched_contours(self):
|
|
for contour_index in range(len(self.fast_contours) -1, -1, -1):
|
|
if self.fast_contours[contour_index].match == None or self.matched_well_contours[contour_index] == None:
|
|
del self.fast_contours[contour_index]
|
|
del self.matched_well_contours[contour_index]
|
|
|
|
def match_contours(self):
|
|
iterations = 0
|
|
while not self.all_contours_matched():
|
|
if iterations >= 5:
|
|
self.delete_unmatched_contours()
|
|
break
|
|
for contour_index, fast_contour in enumerate(self.fast_contours):
|
|
for threshold_index, _ in enumerate(self.well_thresholds):
|
|
for well_contour in self.well_images[threshold_index].contours:
|
|
try:
|
|
well_contour = Contour(well_contour, threshold_index)
|
|
score = fast_contour.get_matching_score(well_contour)
|
|
if score < fast_contour.score:
|
|
if score >= 200:
|
|
continue
|
|
match_index, match = self.existing_contour_match(well_contour)
|
|
if match != None:
|
|
if score < match.score:
|
|
match.unmatch_contour()
|
|
self.matched_well_contours[match_index] = None
|
|
fast_contour.match_contour(well_contour, score)
|
|
self.matched_well_contours[contour_index] = well_contour
|
|
else:
|
|
continue
|
|
else:
|
|
fast_contour.match_contour(well_contour, score)
|
|
self.matched_well_contours[contour_index] = well_contour
|
|
except ContourTooSmallException as e:
|
|
continue
|
|
iterations += 1
|
|
return list(zip(self.matched_well_contours, self.fast_contours))
|
|
|
|
def get_point_descriptors(self, points, image):
|
|
sift = cv2.SIFT_create()
|
|
keypoints = [cv2.KeyPoint(x=float(p[0]), y=float(p[1]), size=20) for p in points]
|
|
keypoints, descriptors = sift.compute(image, keypoints)
|
|
return descriptors
|
|
|
|
def match_points(self, well_contour, fast_contour):
|
|
well_points = well_contour.get_points()
|
|
fast_points = fast_contour.get_points()
|
|
well_descriptors = self.get_point_descriptors(well_points, self.well_images[well_contour.well_image_threshold_index].thresholded)
|
|
fast_descriptors = self.get_point_descriptors(fast_points, self.fast_image.thresholded)
|
|
bf = cv2.BFMatcher()
|
|
matches = bf.match(fast_descriptors, well_descriptors)
|
|
matches = sorted(matches, key=lambda x: x.distance)
|
|
fast_matched_points = [fast_points[m.queryIdx] for m in matches]
|
|
well_matched_points = [well_points[m.trainIdx] for m in matches]
|
|
return well_matched_points, fast_matched_points
|
|
|
|
def alignImages(self):
|
|
matched_contours = self.match_contours()
|
|
well_points_combined = []
|
|
fast_points_combined = []
|
|
for well_contour, fast_contour in matched_contours:
|
|
well_matched_points, fast_matched_points = self.match_points(well_contour, fast_contour)
|
|
well_points_combined.extend(well_matched_points)
|
|
fast_points_combined.extend(fast_matched_points)
|
|
if len(well_points_combined) >= 4:
|
|
well_points = np.float32(well_points_combined)
|
|
fast_points = np.float32(fast_points_combined)
|
|
H, mask = cv2.findHomography(fast_points, well_points, cv2.RANSAC, 5.0)
|
|
warped_image = cv2.warpPerspective(self.fast_image.image, H, (self.image_width, self.image_height))
|
|
else:
|
|
print("Not enough points for homography.")
|
|
return warped_image, H
|
|
|
|
def save_transformation(json_path, right_filename, matrix):
|
|
try:
|
|
with open(json_path, 'r') as f:
|
|
data = json.load(f)
|
|
except (FileNotFoundError, json.JSONDecodeError):
|
|
data = {'matrices': []}
|
|
|
|
found = False
|
|
for item in data['matrices']:
|
|
if item['name'] == right_filename:
|
|
item['matrix'] = matrix.tolist()
|
|
found = True
|
|
break
|
|
|
|
if not found:
|
|
named_matrix = {'name': right_filename, 'matrix': matrix.tolist()}
|
|
data['matrices'].append(named_matrix)
|
|
|
|
with open(json_path, 'w') as f:
|
|
json.dump(data, f, indent=4)
|
|
|
|
def create_mask(matrix):
|
|
matrix = np.array(matrix, dtype=np.float32)
|
|
h, w = 960, 960
|
|
if np.all(matrix == 0):
|
|
return np.ones((h, w), dtype=np.bool_)
|
|
corners = np.array([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]], dtype=np.float32).reshape(-1, 1, 2)
|
|
transformed_corners = cv2.perspectiveTransform(corners, matrix)
|
|
mask = np.zeros((h, w), dtype=np.uint8)
|
|
cv2.fillConvexPoly(mask, np.int32(transformed_corners), 1)
|
|
return mask
|
|
|
|
def delete_unevaluated_images_from_json(transfoms_json_path, eval_images):
|
|
with open(transfoms_json_path, 'r') as json_file:
|
|
data = json.load(json_file)
|
|
|
|
frames_to_keep = []
|
|
for frame in data.get('frames', []):
|
|
file_path = frame.get('file_path', '')
|
|
file_name = os.path.basename(file_path)
|
|
if file_name in eval_images:
|
|
frames_to_keep.append(frame)
|
|
elif 'left_e1' in file_name and file_name.replace('left_e1', 'right_e2') in eval_images:
|
|
frames_to_keep.append(frame)
|
|
data['frames'] = frames_to_keep
|
|
|
|
with open(transfoms_json_path, 'w') as json_file:
|
|
json.dump(data, json_file, indent=4)
|
|
|
|
def change_code_to_eval_all_images(all_images):
|
|
with open('nerfstudio/data/datamanagers/base_datamanager.py', 'r') as file:
|
|
content = file.readlines()
|
|
|
|
new_content = []
|
|
for line in content:
|
|
if 'to eval all images' in line:
|
|
if all_images == True:
|
|
new_content.append(line.replace('# self.eval_dataset', 'self.eval_dataset'))
|
|
else:
|
|
new_content.append(line.replace('self.eval_dataset', '# self.eval_dataset'))
|
|
else:
|
|
new_content.append(line)
|
|
|
|
with open('nerfstudio/data/datamanagers/base_datamanager.py', 'w') as file:
|
|
file.writelines(new_content)
|
|
|
|
if __name__ == '__main__':
|
|
argparser = argparse.ArgumentParser()
|
|
argparser.add_argument("--input_dir", type=str, default="/mnt/data/scene/scene_ns/", help="Nerfstudio directory of the scene")
|
|
argparser.add_argument("--config", type=str, default="outputs/scene_ns/lantern-nerfacto/2024-06-20_142413/config.yml", help="Nerfstudio config file")
|
|
argparser.add_argument("--no_eval", action='store_false', default="True", help="Doesn't run eval if flag is present")
|
|
argparser.add_argument("--no_copy", action='store_false', default="True", help="Doesn't copy the original data if flag is present. Use if already copied.")
|
|
args = argparser.parse_args()
|
|
|
|
if args.no_eval:
|
|
print("Running eval...")
|
|
change_code_to_eval_all_images(True)
|
|
subprocess.run(["ns-eval", "--load-config=" + args.config, "--output-path=output.json", "--render_output_path=" + args.input_dir + "eval_for_alignment"])
|
|
change_code_to_eval_all_images(False)
|
|
|
|
eval_path = os.path.join(args.input_dir, 'eval_for_alignment')
|
|
json_path = os.path.join(args.input_dir, 'alignment_matrices.json')
|
|
transfoms_json_path = os.path.join(args.input_dir, 'transforms.json')
|
|
images_path_aligned = os.path.join(args.input_dir, 'images')
|
|
images_path_unaligned = images_path_aligned + '_unaligned'
|
|
masks_path_aligned = os.path.join(args.input_dir, 'masks')
|
|
masks_path_unaligned = masks_path_aligned + '_unaligned'
|
|
merged_images = os.path.join(images_path_aligned, 'merged')
|
|
|
|
if args.no_copy:
|
|
print("Copying data...")
|
|
subprocess.run(["cp", "-r", images_path_aligned, images_path_unaligned])
|
|
os.makedirs(merged_images, exist_ok=True)
|
|
subprocess.run(["cp", "-r", masks_path_aligned, masks_path_unaligned])
|
|
subprocess.run(["cp", "-r", transfoms_json_path, os.path.join(args.input_dir, 'transforms_unaligned.json')])
|
|
os.makedirs(merged_images, exist_ok=True)
|
|
|
|
print("Aligning images...")
|
|
eval_images = []
|
|
for right_name in os.listdir(images_path_unaligned):
|
|
if right_name.startswith("right") and right_name.endswith(".png"):
|
|
left_eval_image_path = os.path.join(eval_path, right_name.replace('.png', '-img.png'))
|
|
if os.path.exists(left_eval_image_path):
|
|
eval_images.append(right_name)
|
|
left_image = cv2.cvtColor(cv2.imread(left_eval_image_path), cv2.COLOR_BGR2RGB)
|
|
right_path = os.path.join(images_path_unaligned, right_name)
|
|
right_image = cv2.cvtColor(cv2.imread(right_path), cv2.COLOR_BGR2RGB)
|
|
|
|
try:
|
|
align_images = AlignImages(left_image, right_image)
|
|
warped_image, matrix = align_images.alignImages()
|
|
print(f"Aligned: {right_name}")
|
|
except Exception as e:
|
|
warped_image = right_image
|
|
matrix = np.zeros((3, 3))
|
|
print(f"Failed to align: {right_name}")
|
|
|
|
cv2.imwrite(os.path.join(images_path_aligned, right_name), cv2.cvtColor(warped_image, cv2.COLOR_RGB2BGR))
|
|
save_transformation(json_path, right_name, matrix)
|
|
|
|
mask = create_mask(matrix)
|
|
mask_path = os.path.join(masks_path_unaligned, right_name)
|
|
previous_mask = cv2.cvtColor(cv2.imread(mask_path), cv2.COLOR_BGR2GRAY)
|
|
merged_masks = (1 - np.logical_or(1 - mask.astype(np.bool_), 1 - previous_mask.astype(np.bool_))).astype(np.uint8) * 255
|
|
cv2.imwrite(os.path.join(masks_path_aligned, right_name), merged_masks)
|
|
|
|
alpha = 50.0/100.0
|
|
merged_image = cv2.addWeighted(left_image, 1.0-alpha, warped_image, alpha, 0)
|
|
merged_masked_image = cv2.bitwise_or(merged_image, merged_image, mask=merged_masks)
|
|
merged_path = os.path.join(merged_images, right_name + '.merged.png')
|
|
cv2.imwrite(merged_path, cv2.cvtColor(merged_masked_image, cv2.COLOR_RGB2BGR))
|
|
|
|
print("Modifying json...")
|
|
delete_unevaluated_images_from_json(transfoms_json_path, eval_images) |