from PIL import Image import numpy as np import torch, os import sam2point.utils as utils from sam2.build_sam import build_sam2_video_predictor, build_sam2 from sam2.sam2_image_predictor import SAM2ImagePredictor CHECKPOINT = "./checkpoints/sam2_hiera_large.pt" MODELCFG = "sam2_hiera_l.yaml" RESOLUTION = 256 def grid_to_frames(grid, foldpath, args): if not utils.build_fold(foldpath): utils.visualize_per_frame(grid, foldpath=foldpath, resolution=RESOLUTION, args=args) # scan all the JPEG frame names in this directory frame_names = [ p for p in os.listdir(foldpath) if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png"] ] frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) for i in range(len(frame_names)): frame_names[i] = os.path.join(foldpath, frame_names[i]) return frame_names def segment_point(frame_paths, point): sam2_checkpoint = CHECKPOINT model_cfg = MODELCFG predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint) inference_state = predictor.init_state(frame_paths=frame_paths) predictor.reset_state(inference_state) ann_frame_idx = 0 # the frame index we interact with ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers) # Let's add a positive click at (x, y) = (210, 350) to get started points = np.array([point], dtype=np.float32) # for labels, `1` means positive click and `0` means negative click labels = np.array([1], np.int32) _, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box( inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, points=points, labels=labels, ) # run propagation throughout the video and collect the results in a dict video_segments = {} # video_segments contains the per-frame segmentation results for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } masks = [] for out_frame_idx in range(0, len(frame_paths)): for out_obj_id, out_mask in video_segments[out_frame_idx].items(): out_mask = torch.from_numpy(out_mask * 1.0) masks.append(out_mask) masks = torch.cat(masks, dim=0) return masks def segment_box(frame_paths, box, n_frame): sam2_checkpoint = CHECKPOINT model_cfg = MODELCFG predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint) inference_state = predictor.init_state(frame_paths=frame_paths) predictor.reset_state(inference_state) for i in range(n_frame): ann_frame_idx = i # the frame index we interact with ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers) # Let's add a positive click at (x, y) = (210, 350) to get started box = np.array(box, dtype=np.float32) # for labels, `1` means positive click and `0` means negative click _, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box( inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, box=box, ) # run propagation throughout the video and collect the results in a dict video_segments = {} # video_segments contains the per-frame segmentation results for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } masks = [] for out_frame_idx in range(0, len(frame_paths)): for out_obj_id, out_mask in video_segments[out_frame_idx].items(): out_mask = torch.from_numpy(out_mask * 1.0) masks.append(out_mask) masks = torch.cat(masks, dim=0) # print(masks.shape) return masks def segment_mask(frame_paths, point): sam2_checkpoint = CHECKPOINT model_cfg = MODELCFG # generate a mask for one frame, where we use the image predictor sam2_image_model = build_sam2(model_cfg, sam2_checkpoint) image_predictor = SAM2ImagePredictor(sam2_image_model) image = Image.open(frame_paths[0]) image_predictor.set_image(np.array(image.convert("RGB"))) point = np.array([point], dtype=np.float32) label = np.array([1], np.int32) masks, scores, logits = image_predictor.predict(point_coords=point, point_labels=label, multimask_output=True) sorted_ind = np.argsort(scores)[::-1] masks = masks[sorted_ind] # predict the mask for other frames video_predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint) inference_state = video_predictor.init_state(frame_paths=frame_paths) video_predictor.reset_state(inference_state) ann_frame_idx = 0 # the frame index we interact with ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers) mask_prompt = masks[0] video_predictor.add_new_mask(inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, mask=mask_prompt) # run propagation throughout the video and collect the results in a dict video_segments = {} # video_segments contains the per-frame segmentation results for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor.propagate_in_video(inference_state): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } masks = [] for out_frame_idx in range(0, len(frame_paths)): for out_obj_id, out_mask in video_segments[out_frame_idx].items(): out_mask = torch.from_numpy(out_mask * 1.0) masks.append(out_mask) masks = torch.cat(masks, dim=0) return masks, mask_prompt def seg_point(locs, feats, prompt, args): num_voxels = locs.max().astype(int) grid = np.ones((num_voxels + 5, num_voxels+5, num_voxels+5, 3)) # padding locs = locs.astype(int) for v in range(locs.shape[0]): grid[locs[v][0]+2,locs[v][1]+2,locs[v][2]+2] = feats[v] X, Y, Z, _ = grid.shape grid = torch.from_numpy(grid) name_list = ["./tmp/" + args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)] name = '_'.join(name_list) os.makedirs(name + 'frames', exist_ok=True) axis0, axis1, axis2 = name + "frames/x", name + "frames/y", name + "frames/z" grid0, grid1, grid2 = grid.permute(0,3,1,2), grid.permute(1,3,0,2), grid.permute(2,3,0,1) a0_frame_paths = grid_to_frames(grid0, axis0, args) a1_frame_paths = grid_to_frames(grid1, axis1, args) a2_frame_paths = grid_to_frames(grid2, axis2, args) voxel_coords = np.array(prompt) / args.voxel_size + 2 voxel_coords = voxel_coords.astype(int) pixel = voxel_coords * 1.0 / X * RESOLUTION + args.theta * RESOLUTION / X pixel = pixel.astype(int) idx = args.prompt_idx a0_paths_0, a0_paths_1 = a0_frame_paths[:voxel_coords[idx, 0]+1][::-1], a0_frame_paths[voxel_coords[idx, 0]:] a1_paths_0, a1_paths_1 = a1_frame_paths[:voxel_coords[idx, 1]+1][::-1], a1_frame_paths[voxel_coords[idx, 1]:] a2_paths_0, a2_paths_1 = a2_frame_paths[:voxel_coords[idx, 2]+1][::-1], a2_frame_paths[voxel_coords[idx, 2]:] a0_mask_0 = torch.flip(segment_point(a0_paths_0, [pixel[idx, 2], pixel[idx, 1]]), dims=[0]) a0_mask_1 = segment_point(a0_paths_1, [pixel[idx, 2], pixel[idx, 1]])[1:, :, :] a0_mask = torch.cat([a0_mask_0, a0_mask_1], dim=0) a0_mask = torch.nn.functional.interpolate(a0_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a1_mask_0 = torch.flip(segment_point(a1_paths_0, [pixel[idx, 2], pixel[idx, 0]]), dims=[0]) a1_mask_1 = segment_point(a1_paths_1, [pixel[idx, 2], pixel[idx, 0]])[1:, :, :] a1_mask = torch.cat([a1_mask_0, a1_mask_1], dim=0) a1_mask = torch.nn.functional.interpolate(a1_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a2_mask_0 = torch.flip(segment_point(a2_paths_0, [pixel[idx, 1], pixel[idx, 0]]), dims=[0]) a2_mask_1 = segment_point(a2_paths_1, [pixel[idx, 1], pixel[idx, 0]])[1:, :, :] a2_mask = torch.cat([a2_mask_0, a2_mask_1], dim=0) a2_mask = torch.nn.functional.interpolate(a2_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a0_mask, a1_mask, a2_mask = a0_mask.transpose(0, 1), a1_mask.transpose(0, 1), a2_mask.transpose(0, 1) # utils.visualize_frame_with_mask(grid0, grid1, grid2, a0_mask, a1_mask, a2_mask, voxel_coords[idx], resolution=RESOLUTION) mask = a0_mask.permute(0, 2, 3, 1) + a1_mask.permute(2, 0, 3, 1) + a2_mask.permute(2, 3, 0, 1) mask = (mask > 1.5).squeeze()[2:, 2:, 2:] return mask def seg_box(locs, feats, prompt, args): num_voxels = locs.max().astype(int) grid = np.ones((num_voxels + 5, num_voxels+5, num_voxels+5, 3)) # padding locs = locs.astype(int) for v in range(locs.shape[0]): grid[locs[v][0]+2,locs[v][1]+2,locs[v][2]+2] = feats[v] X, Y, Z, _ = grid.shape grid = torch.from_numpy(grid) name_list = ["./tmp/" + args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)] name = '_'.join(name_list) os.makedirs(name + 'frames', exist_ok=True) axis0, axis1, axis2 = name + "frames/x", name + "frames/y", name + "frames/z" grid0, grid1, grid2 = grid.permute(0,3,1,2), grid.permute(1,3,0,2), grid.permute(2,3,0,1) a0_frame_paths = grid_to_frames(grid0, axis0, args) a1_frame_paths = grid_to_frames(grid1, axis1, args) a2_frame_paths = grid_to_frames(grid2, axis2, args) point_prompts = np.array(prompt) voxel_coords = point_prompts / args.voxel_size + 2 voxel_coords = voxel_coords.astype(int) pixel = voxel_coords * 1.0 / X * RESOLUTION + args.theta * RESOLUTION / X pixel = pixel.astype(int) idx = args.prompt_idx a0_paths_0, a0_paths_1 = a0_frame_paths[:voxel_coords[idx, 3]+1][::-1], a0_frame_paths[voxel_coords[idx, 0]:] a1_paths_0, a1_paths_1 = a1_frame_paths[:voxel_coords[idx, 4]+1][::-1], a1_frame_paths[voxel_coords[idx, 1]:] a2_paths_0, a2_paths_1 = a2_frame_paths[:voxel_coords[idx, 5]+1][::-1], a2_frame_paths[voxel_coords[idx, 2]:] frame_num0 = voxel_coords[idx, 3] - voxel_coords[idx, 0] end0, start0 = len(a0_paths_0) - int(frame_num0 / 2), int(frame_num0 / 2) a0_mask_0 = torch.flip(segment_box(a0_paths_0, [pixel[idx, 2], pixel[idx, 1], pixel[idx, 5], pixel[idx, 4]], frame_num0), dims=[0])[:end0] a0_mask_1 = segment_box(a0_paths_1, [pixel[idx, 2], pixel[idx, 1], pixel[idx, 5], pixel[idx, 4]], frame_num0)[start0:] a0_mask = torch.cat([a0_mask_0, a0_mask_1], dim=0) a0_mask = torch.nn.functional.interpolate(a0_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) frame_num1 = voxel_coords[idx, 4] - voxel_coords[idx, 1] end1, start1 = len(a1_paths_0) - int(frame_num1 / 2), int(frame_num1 / 2) a1_mask_0 = torch.flip(segment_box(a1_paths_0, [pixel[idx, 2], pixel[idx, 0], pixel[idx, 5], pixel[idx, 3]], frame_num1), dims=[0])[:end1] a1_mask_1 = segment_box(a1_paths_1, [pixel[idx, 2], pixel[idx, 0], pixel[idx, 5], pixel[idx, 3]], frame_num1)[start1:] a1_mask = torch.cat([a1_mask_0, a1_mask_1], dim=0) a1_mask = torch.nn.functional.interpolate(a1_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) frame_num2 = voxel_coords[idx, 5] - voxel_coords[idx, 2] end2, start2 = len(a2_paths_0) - int(frame_num2 / 2), int(frame_num2 / 2) a2_mask_0 = torch.flip(segment_box(a2_paths_0, [pixel[idx, 1], pixel[idx, 0], pixel[idx, 4], pixel[idx, 3]], frame_num2), dims=[0])[:end2] a2_mask_1 = segment_box(a2_paths_1, [pixel[idx, 1], pixel[idx, 0], pixel[idx, 4], pixel[idx, 3]], frame_num2)[start2:] a2_mask = torch.cat([a2_mask_0, a2_mask_1], dim=0) a2_mask = torch.nn.functional.interpolate(a2_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a0_mask, a1_mask, a2_mask = a0_mask.transpose(0, 1), a1_mask.transpose(0, 1), a2_mask.transpose(0, 1) # utils.visualize_frame_with_mask(grid0, grid1, grid2, a0_mask, a1_mask, a2_mask, voxel_coords[idx], resolution=RESOLUTION) mask = a0_mask.permute(0, 2, 3, 1) + a1_mask.permute(2, 0, 3, 1) + a2_mask.permute(2, 3, 0, 1) mask = (mask > 1.5).squeeze()[2:, 2:, 2:] return mask def seg_mask(locs, feats, prompt, args): num_voxels = locs.max().astype(int) grid = np.ones((num_voxels + 5, num_voxels+5, num_voxels+5, 3)) # padding locs = locs.astype(int) for v in range(locs.shape[0]): grid[locs[v][0]+2,locs[v][1]+2,locs[v][2]+2] = feats[v] X, Y, Z, _ = grid.shape grid = torch.from_numpy(grid) name_list = ["./tmp/" + args.dataset, "sample" + str(args.sample_idx), args.prompt_type + "-prompt" + str(args.prompt_idx)] name = '_'.join(name_list) os.makedirs(name + 'frames', exist_ok=True) axis0, axis1, axis2 = name + "frames/x", name + "frames/y", name + "frames/z" grid0, grid1, grid2 = grid.permute(0,3,1,2), grid.permute(1,3,0,2), grid.permute(2,3,0,1) a0_frame_paths = grid_to_frames(grid0, axis0, args) a1_frame_paths = grid_to_frames(grid1, axis1, args) a2_frame_paths = grid_to_frames(grid2, axis2, args) point_prompts = np.array(prompt) voxel_coords = point_prompts / args.voxel_size + 2 voxel_coords = voxel_coords.astype(int) pixel = voxel_coords * 1.0 / X * RESOLUTION + args.theta * RESOLUTION / X pixel = pixel.astype(int) idx = args.prompt_idx a0_paths_0, a0_paths_1 = a0_frame_paths[:voxel_coords[idx, 0]+1][::-1], a0_frame_paths[voxel_coords[idx, 0]:] a1_paths_0, a1_paths_1 = a1_frame_paths[:voxel_coords[idx, 1]+1][::-1], a1_frame_paths[voxel_coords[idx, 1]:] a2_paths_0, a2_paths_1 = a2_frame_paths[:voxel_coords[idx, 2]+1][::-1], a2_frame_paths[voxel_coords[idx, 2]:] a0_mask_0, a0_prompt = segment_mask(a0_paths_0, [pixel[idx, 2], pixel[idx, 1]]) a0_mask_0 = torch.flip(a0_mask_0, dims=[0]) a0_mask_1, _ = segment_mask(a0_paths_1, [pixel[idx, 2], pixel[idx, 1]]) a0_mask_1 = a0_mask_1[1:, :, :] a0_mask = torch.cat([a0_mask_0, a0_mask_1], dim=0) a0_prompt_mask = a0_mask * 0 a0_prompt_mask[voxel_coords[idx, 0]] = torch.from_numpy(a0_prompt) a0_mask = torch.nn.functional.interpolate(a0_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a0_prompt_mask = torch.nn.functional.interpolate(a0_prompt_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a1_mask_0, a1_prompt = segment_mask(a1_paths_0, [pixel[idx, 2], pixel[idx, 0]]) a1_mask_0 = torch.flip(a1_mask_0, dims=[0]) a1_mask_1, _ = segment_mask(a1_paths_1, [pixel[idx, 2], pixel[idx, 0]]) a1_mask_1 = a1_mask_1[1:, :, :] a1_mask = torch.cat([a1_mask_0, a1_mask_1], dim=0) a1_prompt_mask = a1_mask * 0 a1_prompt_mask[voxel_coords[idx, 1]] = torch.from_numpy(a1_prompt) a1_mask = torch.nn.functional.interpolate(a1_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a1_prompt_mask = torch.nn.functional.interpolate(a1_prompt_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a2_mask_0, a2_prompt = segment_mask(a2_paths_0, [pixel[idx, 1], pixel[idx, 0]]) a2_mask_0 = torch.flip(a2_mask_0, dims=[0]) a2_mask_1, _ = segment_mask(a2_paths_1, [pixel[idx, 1], pixel[idx, 0]]) a2_mask_1 = a2_mask_1[1:, :, :] a2_mask = torch.cat([a2_mask_0, a2_mask_1], dim=0) a2_prompt_mask = a2_mask * 0 a2_prompt_mask[voxel_coords[idx, 2]] = torch.from_numpy(a2_prompt) a2_mask = torch.nn.functional.interpolate(a2_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a2_prompt_mask = torch.nn.functional.interpolate(a2_prompt_mask.unsqueeze(0).unsqueeze(0), size=(X, X, X), mode='trilinear').squeeze(0) a0_mask, a1_mask, a2_mask = a0_mask.transpose(0, 1), a1_mask.transpose(0, 1), a2_mask.transpose(0, 1) utils.visualize_frame_with_mask(grid0, grid1, grid2, a0_mask, a1_mask, a2_mask, voxel_coords[idx], resolution=RESOLUTION, name=name, args=args) a0_prompt_mask, a1_prompt_mask, a2_prompt_mask = a0_prompt_mask.transpose(0, 1), a1_prompt_mask.transpose(0, 1), a2_prompt_mask.transpose(0, 1) mask = a0_mask.permute(0, 2, 3, 1) + a1_mask.permute(2, 0, 3, 1) + a2_mask.permute(2, 3, 0, 1) mask = (mask > 1.5).squeeze()[2:, 2:, 2:] prompt_mask = a0_prompt_mask.permute(0, 2, 3, 1) + a1_prompt_mask.permute(2, 0, 3, 1) + a2_prompt_mask.permute(2, 3, 0, 1) prompt_mask = (prompt_mask > 0.5).squeeze()[2:, 2:, 2:] return mask, prompt_mask