diff --git a/lantern_scripts/mask_humans2.py b/lantern_scripts/mask_humans2.py index dc1d4db1a..02eb22743 100644 --- a/lantern_scripts/mask_humans2.py +++ b/lantern_scripts/mask_humans2.py @@ -14,6 +14,8 @@ from PIL import Image # import some common detectron2 utilities from tqdm import tqdm +from envmap import EnvironmentMap, rotation_matrix + FRAME_SUBFOLDERS_CONFIG = [ { @@ -74,16 +76,29 @@ if __name__ == '__main__': original_image_pil = Image.open(os.path.join(args.input_dir, folder['name'], file)).convert("RGB") original_image_np = np.array(original_image_pil) + original_image = EnvironmentMap(os.path.join(args.input_dir, folder['name'], file), 'latlong') + dcm = rotation_matrix(azimuth=0, elevation=0, roll=np.pi/2) + original_image.rotate(dcm) + image_np_rot = (original_image.data * 255).astype(np.uint8) + if folder['mask_person']: text_prompt = "person" # 90 degrees rotat - image_np_rot = np.rot90(original_image_np, k=folder['flip_k']) + # image_np_rot = original_image_np # np.rot90(original_image_np, k=folder['flip_k']) masks, boxes, phrases, logits = model.predict(Image.fromarray(image_np_rot), text_prompt) masks = masks.detach().cpu().numpy() - masks = np.rot90(masks, k=-folder['flip_k'], axes=(1, 2)) - - # aggregate masks - masks = np.max(masks, axis=0) + if masks.size != 0: + # masks = np.rot90(masks, k=-folder['flip_k'], axes=(1, 2)) + + mask_skylibs = EnvironmentMap(np.transpose(masks, (1, 2, 0)).astype(np.uint8), 'latlong') + inv_dcm = rotation_matrix(azimuth=0, elevation=0, roll=-np.pi/2) + mask_skylibs.rotate(inv_dcm) + masks = mask_skylibs.data + masks = np.transpose(masks, (2, 0, 1)) + # aggregate masks + masks = np.max(masks, axis=0) + else: + masks = np.zeros((original_image_pil.height, original_image_pil.width)) else: masks = np.zeros((original_image_pil.height, original_image_pil.width)) diff --git a/nerfstudio/process_data/GT_HDR_to_nerfstudio.py b/nerfstudio/process_data/GT_HDR_to_nerfstudio.py index 6c6ce5009..d8d18c1a9 100644 --- a/nerfstudio/process_data/GT_HDR_to_nerfstudio.py +++ b/nerfstudio/process_data/GT_HDR_to_nerfstudio.py @@ -107,20 +107,17 @@ class GTHDRoNeRFStudioDataset(BaseConverterToNerfstudioDataset): scale_factor = dataset_transform['scale'] normalization_mat = np.eye(4) normalization_mat[:3, :] = np.array(dataset_transform['transform']) - + # Convert OpenSFM coordinates to the NeRFStudio camera_dict = {} camera_dict["fps"] = 1 - camera_key = 'v2 unknown unknown -1 -1 perspective 0' - # camera_key = 'v2 unknown unknown 3840 1920 perspective 0.0' - # camera_key = 'v2 unknown unknown 6720 3360 perspective 0.0' camera_dict["camera_type"] = "equirectangular" - camera_dict["render_width"] = openSFM_reconstruction[0]["cameras"][camera_key]['width'] - camera_dict["render_height"] = openSFM_reconstruction[0]["cameras"][camera_key]['height'] + camera_dict["render_width"] = 3840 + camera_dict["render_height"] = 1920 camera_dict["smoothness_value"] = 0 camera_dict["is_cycle"] = False camera_dict["camera_path"] = [] - + frame_length = 0 for idx, fname in enumerate([x for x in frames_names if x.startswith('GT')]): print(fname) diff --git a/stick_masks/left_mask.png b/stick_masks/left_mask.png index 03ba7fbdc..8905028b2 100644 Binary files a/stick_masks/left_mask.png and b/stick_masks/left_mask.png differ diff --git a/stick_masks/right_mask.png b/stick_masks/right_mask.png index b447835f5..f35aa374c 100644 Binary files a/stick_masks/right_mask.png and b/stick_masks/right_mask.png differ