diff --git a/.hydra/config.yaml b/.hydra/config.yaml index 73f0cd9..d9b8bfd 100644 --- a/.hydra/config.yaml +++ b/.hydra/config.yaml @@ -1,27 +1,23 @@ work_dir: ${hydra:runtime.cwd} data_dir: ${work_dir}/data print_config: true -type: inference -task_name: demo -num_leaf: 8 -suffix: '' -save_demo: false -save_wis3d: false -use_tracking: false -model: - onepose_model_path: ${work_dir}/data/models/checkpoints/onepose/GATsSPG.ckpt - extractor_model_path: ${work_dir}/data/models/extractors/SuperPoint/superpoint_v1.pth - match_model_path: ${work_dir}/data/models/matchers/SuperGlue/superglue_outdoor.pth -scan_data_dir: ${data_dir}/onepose_datasets/val_data -sfm_model_dir: ${data_dir}/sfm_model +type: sfm +redo: true +hololens: false +scan_data_dir: ${work_dir}/data/onepose_datasets/val_data +dataset: + max_num_kp3d: 2500 + max_num_kp2d: 1000 + data_dir: + - ${scan_data_dir}/0620-dinosaurcup-bottle dinosaurcup-1 + outputs_dir: ${work_dir}/data/sfm_model/{} network: detection: superpoint + detection_model_path: ${work_dir}/data/models/extractors/SuperPoint/superpoint_v1.pth matching: superglue -max_num_kp3d: 2500 -input: - data_dirs: /Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/onepose_datasets/tiger - tiger-test - sfm_model_dirs: /Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/sfm_model/tiger -output: - vis_dir: ${work_dir}/runs/vis/demo - eval_dir: ${work_dir}/runs/eval/demo + matching_model_path: ${work_dir}/data/models/matchers/SuperGlue/superglue_outdoor.pth +sfm: + down_ratio: 5 + covis_num: 10 + rotation_thresh: 50 +disable_lightning_logs: true diff --git a/.hydra/hydra.yaml b/.hydra/hydra.yaml index 3f952ed..1c6a8f8 100644 --- a/.hydra/hydra.yaml +++ b/.hydra/hydra.yaml @@ -109,15 +109,10 @@ hydra: overrides: hydra: [] task: - - +experiment=test_demo - - use_tracking=False - - input.data_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/onepose_datasets/tiger - tiger-test - - input.sfm_model_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/sfm_model/tiger + - +preprocess=sfm_spp_spg_val.yaml job: - name: inference_demo - override_dirname: +experiment=test_demo,input.data_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/onepose_datasets/tiger - tiger-test,input.sfm_model_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/sfm_model/tiger,use_tracking=False + name: run + override_dirname: +preprocess=sfm_spp_spg_val.yaml id: ??? num: ??? config_name: config.yaml @@ -130,19 +125,19 @@ hydra: exclude_keys: [] runtime: version: 1.1.1 - cwd: /Users/PELLERITO/Desktop/mixed_reality_code/OnePose + cwd: /Users/diego/Desktop/Escritorio_MacBook_Pro_de_Diego/ETH/Third_Semester/Mixed_Reality/Real_2/ZeroShotPoseEstimation config_sources: - path: hydra.conf schema: pkg provider: hydra - - path: /Users/PELLERITO/Desktop/mixed_reality_code/OnePose/configs + - path: /Users/diego/Desktop/Escritorio_MacBook_Pro_de_Diego/ETH/Third_Semester/Mixed_Reality/Real_2/ZeroShotPoseEstimation/configs schema: file provider: main - path: '' schema: structured provider: schema choices: - experiment: test_demo + preprocess: sfm_spp_spg_val.yaml logger: null callbacks: null datamodule: null diff --git a/.hydra/overrides.yaml b/.hydra/overrides.yaml index 31687cf..4dcc863 100644 --- a/.hydra/overrides.yaml +++ b/.hydra/overrides.yaml @@ -1,5 +1 @@ -- +experiment=test_demo -- use_tracking=False -- input.data_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/onepose_datasets/tiger - tiger-test -- input.sfm_model_dirs=/Users/PELLERITO/Desktop/mixed_reality_code/OnePose/data/sfm_model/tiger +- +preprocess=sfm_spp_spg_val.yaml diff --git a/Output/tiger/ellipsoids.png b/Output/tiger/ellipsoids.png new file mode 100644 index 0000000..f2932c6 Binary files /dev/null and b/Output/tiger/ellipsoids.png differ diff --git a/configs/preprocess/sfm_spp_spg_val.yaml b/configs/preprocess/sfm_spp_spg_val.yaml index ef0f599..f2b9021 100644 --- a/configs/preprocess/sfm_spp_spg_val.yaml +++ b/configs/preprocess/sfm_spp_spg_val.yaml @@ -3,18 +3,17 @@ type: sfm work_dir: ${hydra:runtime.cwd} redo: True -hololens: True +hololens: False # scan_data_dir: ${work_dir}/data/onepose_datasets/val_data -scan_data_dir: ${work_dir}/data/costum_datasets/test - +scan_data_dir: ${work_dir}/data/onepose_datasets/val_data dataset: max_num_kp3d: 2500 max_num_kp2d: 1000 data_dir: # - ${scan_data_dir}/0606-tiger-others tiger-2 - - ${scan_data_dir}/demo_capture demo-1 + - ${scan_data_dir}/0620-dinosaurcup-bottle dinosaurcup-1 outputs_dir: ${work_dir}/data/sfm_model/{} diff --git a/feature_matching_object_detector.py b/feature_matching_object_detector.py index 01f5cd5..af767ab 100644 --- a/feature_matching_object_detector.py +++ b/feature_matching_object_detector.py @@ -149,9 +149,6 @@ def inference_core(cfg, data_root, seq_dir, sfm_model_dir): # Begin Object detection: for id, data in enumerate(tqdm(loader)): - #for data in (tqdm(loader)): - - #for id, data in enumerate((loader)): img_path = data["path"][0] inp = data["image"].to(device) diff --git a/run.py b/run.py index 8b39ef2..d052ac4 100644 --- a/run.py +++ b/run.py @@ -121,11 +121,12 @@ def sfm(cfg): for sub_dir in sub_dirs: seq_dir = osp.join(root_dir, sub_dir) img_paths += glob.glob(str(Path(seq_dir)) + "/color/*.png", recursive=True) - full_res_img_paths += glob.glob(str(Path(seq_dir)) + "/color_full/*.png", recursive=True) - poses_paths += glob.glob(str(Path(seq_dir)) + "/poses/*.txt", recursive=True) - intrinsics_path = str(Path(seq_dir)) + "/intrinsics.txt" + full_res_img_paths += glob.glob(str(Path(seq_dir)) + "/color_full/*.jpg", recursive=True) poses_paths += glob.glob(str(Path(seq_dir)) + "/poses/*.txt", recursive=True) intrinsics_path = str(Path(seq_dir)) + "/intrinsics.txt" + + #poses_paths += glob.glob(str(Path(seq_dir)) + "/poses/*.txt", recursive=True) + #intrinsics_path = str(Path(seq_dir)) + "/intrinsics.txt" # Choose less images from the list to build the sfm model @@ -149,14 +150,20 @@ def sfm(cfg): poses_paths=poses_paths, data_root=root_dir, seq_dir = seq_dir, - compute_on_GPU="cuda", + compute_on_GPU="cpu", step=1, - hololens=cfg.hololens + hololens=cfg.hololens, + filter = False, #DSM Filter ) + # Begin SfM and postprocess: + print("Beginning of SFM") sfm_core(cfg, down_img_lists, outputs_dir_root) + print("SfM done") + print("Beginning of Post-process") postprocess(cfg, down_img_lists, root_dir, outputs_dir_root) + print("Post-process done") def sfm_core(cfg, img_lists, outputs_dir_root): @@ -180,6 +187,7 @@ def sfm_core(cfg, img_lists, outputs_dir_root): empty_dir = osp.join(outputs_dir, "sfm_empty") deep_sfm_dir = osp.join(outputs_dir, "sfm_ws") + if cfg.redo: os.system(f"rm -rf {outputs_dir}") Path(outputs_dir).mkdir(exist_ok=True, parents=True) diff --git a/src/bbox_3D_estimation/utils.py b/src/bbox_3D_estimation/utils.py index 196aa3c..cac8cb5 100644 --- a/src/bbox_3D_estimation/utils.py +++ b/src/bbox_3D_estimation/utils.py @@ -41,6 +41,7 @@ def detect_3D_box(self): estQs = compute_estimates(self.bboxes, self.K, self.poses, self.visibility) centre, axes, R = dual_quadric_to_ellipsoid_parameters(estQs[object_idx]) + print("Axes = ", axes) # Possible coordinates mins = [-ax for (ax) in axes] maxs = [ax for (ax) in axes] @@ -58,31 +59,17 @@ def detect_3D_box(self): self.points = points self.centre = centre self.R = R + self.estQs = estQs # Transformation to have coordinates centered in the bounding box (and aligned with it) M = np.empty((4, 4)) M[:3, :3] = R M[:3, 3] = centre M[3, :] = [0, 0, 0, 1] - # print(M) self.M = np.linalg.inv(M) - # gt_p = np.loadtxt(f"data/onepose_datasets/val_data/0606-tiger-others/box3d_corners_GT.txt") - - # plot_3D_scene( - # estQs=estQs, - # gtQs=gt_p, - # Ms_t=self.poses, - # dataset="tiger", - # save_output_images=False, - # points=points, - # GT_points=gt_p - # ) - # plt.show() - - def save_3D_box(self, data_root): np.savetxt(data_root + "/box3d_corners.txt", self.points, delimiter=" ") @@ -105,6 +92,18 @@ def save_poses(self, seq_dir): def save_dimensions(self, data_root): np.savetxt(data_root + "/box3d_dimensions.txt", self.axes, delimiter=" ") + def plot_3D_bb(self, poses, GT_points=None): + plot_3D_scene( + estQs=self.estQs, + gtQs=None, + Ms_t=poses, + dataset="tiger", + save_output_images=True, + points=self.points, + GT_points=GT_points + ) + plt.show() + def predict_3D_bboxes( full_res_img_paths, @@ -112,26 +111,54 @@ def predict_3D_bboxes( poses_paths, data_root, seq_dir, + filter, step=1, - downscale_factor=0.3, + downscale_factor=0.6, compute_on_GPU="cpu", hololens=False + ): full_res_img_paths = sort_path_list(full_res_img_paths) poses_paths = sort_path_list(poses_paths) _K, _ = data_utils.get_K(intrisics_path) DetectorBox3D = Detector3D(_K) - BboxPredictor = UnsupBbox(downscale_factor=downscale_factor, device=compute_on_GPU) + BboxPredictor = UnsupBbox(downscale_factor=downscale_factor, device=compute_on_GPU, filter = filter) + for id, img_path in enumerate(tqdm(full_res_img_paths)): if id % step == 0 or id == 0: + image = cv2.imread(str(img_path)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + poses = read_list_poses([poses_paths[id]], hololens=hololens) poses_orig = read_list_poses_orig([poses_paths[id]]) - bbox_orig_res = BboxPredictor.infer_2d_bbox(image=image, K=_K) + + + ''' + ##--------Plotting the RGB image--------- + image_2 = image.copy() + limits = bbox_orig_res + start_point = (int(limits[0]), int(limits[1])) + end_point = (int(limits[2]), int(limits[3])) + + # Blue color in BGR + color = (255, 0, 0) + # Line thickness of 2 px + thickness = 2 + + # Draw a rectangle with blue line borders of thickness of 2 px + image_2 = cv2.rectangle(image_2, start_point, end_point, color, thickness) + + cv2.imshow("Image", image_2) + #cv2.imwrite('/Users/diego/Desktop/Escritorio_MacBook_Pro_de_Diego/ETH/Third_Semester/Mixed_Reality/Real_2/ZeroShotPoseEstimation/data/onepose_datasets/val_data/0620-dinosaurcup-bottle/dinosaurcup-1/detection_2/'+str(id)+ '.jpg', image_2) + cv2.waitKey(1) + ''' + + + #1 is for the bbox without a filter, 0 is with filter DetectorBox3D.add_view(bbox_orig_res, poses, poses_orig) DetectorBox3D.detect_3D_box() @@ -140,6 +167,20 @@ def predict_3D_bboxes( DetectorBox3D.save_poses(seq_dir) DetectorBox3D.save_dimensions(data_root) + poses_t=None + for id, img_path in enumerate(tqdm(full_res_img_paths)): + poses = read_list_poses([poses_paths[id]], hololens=False) + if poses_t is None: + poses_t = poses + else: + poses_t = np.vstack((poses_t, poses)) + + #DetectorBox3D.plot_3D_bb(poses_t) + #print("-------Predict 3D bboxes finished-------") + + GT_points = np.loadtxt(data_root + "/box3d_corners_GT.txt") + + DetectorBox3D.plot_3D_bb(poses_t, GT_points=GT_points) def sort_path_list(path_list): files = {int(Path(file).stem): file for file in path_list} @@ -477,6 +518,7 @@ def estimate_one_ellipsoid(Ps_t, Cs): i += 1 print(i) + _, _, V = np.linalg.svd(M) w = V[-1, :] # V is transposed respect to Matlab, so we take the last row. diff --git a/src/deep_spectral_method/detection_2D_utils.py b/src/deep_spectral_method/detection_2D_utils.py index 3a11de4..1ca89e6 100644 --- a/src/deep_spectral_method/detection_2D_utils.py +++ b/src/deep_spectral_method/detection_2D_utils.py @@ -3,16 +3,23 @@ import torch import numpy as np import cv2 +import lmfit +import matplotlib.pyplot as plt # TODO: extend this class to load dino model on the constructor class UnsupBbox: - def __init__(self, downscale_factor=0.3, device="cpu") -> None: + def __init__(self, downscale_factor=0.3, device="cpu", filter = False) -> None: + #Filter needs a bigger image + if filter: + downscale_factor=0.6 self.model_name = "dino_vits16" self.num_workers = 0 # decrease this if out_of_memory error self.downscale_factor = downscale_factor self.device = device + self.fitting_model = lmfit.models.Gaussian2dModel() #Model to fit the gaussian curve + self.filter = filter if device=="cpu": self.on_GPU = False @@ -31,6 +38,7 @@ def __init__(self, downscale_factor=0.3, device="cpu") -> None: self.patch_size, self.num_heads, ) = utils.get_model(self.model_name) + self.model = self.model.to(device) def downscale_image(self, image): @@ -56,14 +64,44 @@ def infer_2d_bbox(self, image, K): K=4, data_dict=feature_dict, device=self.device ) - # small Segmentation - segmap = extract.extract_single_region_segmentations( - feature_dict=feature_dict, eigs_dict=eigs_dict - ) + eigs_dict_gaussian = eigs_dict.copy() - # Bounding boxes - bbox = extract.extract_bboxes(feature_dict=feature_dict, segmap=segmap) - bbox_orig_res = ( - np.array(bbox["bboxes_original_resolution"][0]) / self.downscale_factor - ) - return bbox_orig_res + #Use of the filter, by default is set to False + if self.filter: + # small segmentation, use of fitting curve + segmap = extract.gaussian_fitting( + feature_dict=feature_dict, + eigs_dict=eigs_dict_gaussian, + fitting_model=self.fitting_model + ) + + bbox = extract.extract_bboxes( + feature_dict=feature_dict, + segmap=segmap, + ) + + #If we cannot find a solution with the filter then use the original bbox + if not bbox['bboxes']: + # small Segmentation + segmap = extract.extract_single_region_segmentations( + feature_dict=feature_dict, eigs_dict=eigs_dict + ) + # Bounding boxes + bbox = extract.extract_bboxes(feature_dict=feature_dict, segmap=segmap) + + else: + + # small Segmentation + segmap = extract.extract_single_region_segmentations( + feature_dict=feature_dict, eigs_dict=eigs_dict + ) + # Bounding boxes + bbox = extract.extract_bboxes(feature_dict=feature_dict, segmap=segmap) + + self.bbox_orig_res = ( + np.array(bbox["bboxes_original_resolution"][0]) / self.downscale_factor + ) + return self.bbox_orig_res + + def save_2d_bbox(self, file_path): + np.savetxt(file_path, self.bbox_orig_res, delimiter=" ") diff --git a/src/deep_spectral_method/extract/extract.py b/src/deep_spectral_method/extract/extract.py index d144006..7341fa7 100644 --- a/src/deep_spectral_method/extract/extract.py +++ b/src/deep_spectral_method/extract/extract.py @@ -17,10 +17,80 @@ from scipy.sparse.linalg import eigsh from . import extract_utils as utils - +import lmfit # import extract_utils as utils +def gaussian_fitting( + feature_dict: dict, + eigs_dict: dict, + fitting_model:dict, +): + flag = False + data_dict = feature_dict + data_dict.update(eigs_dict) + # Sizes + B, C, H, W, P, H_patch, W_patch, H_pad, W_pad = utils.get_image_sizes(data_dict) + + Xin, Yin = np.mgrid[0:H_patch, 0:W_patch] + # Eigenvector + eigenvector = data_dict["eigenvectors"][ + 1 + ].numpy().copy() + + params = fitting_model.guess(eigenvector, Xin.flatten(), Yin.flatten()) + params['sigmax'].set(min = 0, max = W_patch/2) + params['sigmay'].set(min = 0, max = H_patch/2) + + #This happens when the eigenvector is flipped + if (np.max(eigenvector) < 0.01): + #Get the lower 10% + ei = np.sort(eigenvector)[int(0.10 * W_patch * H_patch)] + eigenvector[eigenvector > ei] = 0. + + eigs = eigenvector.copy().reshape(H_patch, W_patch) + graph_2 = (np.array(eigs) / np.min((eigs))) + + mult_x = graph_2 * Xin + mult_y = graph_2 * Yin + + xsum = np.mean(mult_x[eigs != 0.]) + ysum = np.mean(mult_y[eigs != 0.]) + + params['centerx'].set(value = xsum, min = 0, max = W_patch) + params['centery'].set(value = ysum, min = 0, max = H_patch) + + params['amplitude'].set(np.min(eigenvector)) + flag = True + + result = fitting_model.fit(eigenvector.flatten(), x=Xin.flatten(), y=Yin.flatten(), params=params, weights=1/np.sqrt(eigenvector+1)) + fit = fitting_model.func(Xin, Yin, **result.best_values) + + #Flip the fitted curve + if flag: + fit *= -1 + + fit += -np.min(fit) + fit /= np.max(fit) + + #Check the eigenvectors + ''' + fit_2 = fit.copy() + eigs = eigenvector.reshape(H_patch, W_patch) + plt.imshow(eigs) + plt.show() + + cv2.namedWindow("Fitting",cv2.WINDOW_NORMAL) + cv2.imshow("Fitting", fit_2) + cv2.resizeWindow("Fitting", 600,600) + cv2.waitKey(1) + ''' + + #Segmentation mask + segmap = (fit > 0.3) + + return segmap + def extract_features( model: dict, patch_size: int, diff --git a/src/sfm/extract_features.py b/src/sfm/extract_features.py index d566d56..7f92f39 100644 --- a/src/sfm/extract_features.py +++ b/src/sfm/extract_features.py @@ -38,7 +38,7 @@ def spp(img_lists, feature_out, cfg): load_network(model, cfg.network.detection_model_path, force=True) dataset = NormalizedDataset(img_lists, conf["preprocessing"]) - loader = DataLoader(dataset, num_workers=1) + loader = DataLoader(dataset, num_workers=0) feature_file = h5py.File(feature_out, "w") logging.info(f"Exporting features to {feature_out}")