yslan/ObjCtrl-2.5D
10
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3 4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import os8import warnings9from threading import Thread10 11import numpy as np12import torch13from PIL import Image14from tqdm import tqdm15 16 17def get_sdpa_settings():18 if torch.cuda.is_available():19 old_gpu = torch.cuda.get_device_properties(0).major < 720 # only use Flash Attention on Ampere (8.0) or newer GPUs21 use_flash_attn = torch.cuda.get_device_properties(0).major >= 822 if not use_flash_attn:23 warnings.warn(24 "Flash Attention is disabled as it requires a GPU with Ampere (8.0) CUDA capability.",25 category=UserWarning,26 stacklevel=2,27 )28 # keep math kernel for PyTorch versions before 2.2 (Flash Attention v2 is only29 # available on PyTorch 2.2+, while Flash Attention v1 cannot handle all cases)30 pytorch_version = tuple(int(v) for v in torch.__version__.split(".")[:2])31 if pytorch_version < (2, 2):32 warnings.warn(33 f"You are using PyTorch {torch.__version__} without Flash Attention v2 support. "34 "Consider upgrading to PyTorch 2.2+ for Flash Attention v2 (which could be faster).",35 category=UserWarning,36 stacklevel=2,37 )38 math_kernel_on = pytorch_version < (2, 2) or not use_flash_attn39 else:40 old_gpu = True41 use_flash_attn = False42 math_kernel_on = True43 44 return old_gpu, use_flash_attn, math_kernel_on45 46 47def get_connected_components(mask):48 """49 Get the connected components (8-connectivity) of binary masks of shape (N, 1, H, W).50 51 Inputs:52 - mask: A binary mask tensor of shape (N, 1, H, W), where 1 is foreground and 0 is53 background.54 55 Outputs:56 - labels: A tensor of shape (N, 1, H, W) containing the connected component labels57 for foreground pixels and 0 for background pixels.58 - counts: A tensor of shape (N, 1, H, W) containing the area of the connected59 components for foreground pixels and 0 for background pixels.60 """61 from sam2 import _C62 63 return _C.get_connected_componnets(mask.to(torch.uint8).contiguous())64 65 66def mask_to_box(masks: torch.Tensor):67 """68 compute bounding box given an input mask69 70 Inputs:71 - masks: [B, 1, H, W] masks, dtype=torch.Tensor72 73 Returns:74 - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor75 """76 B, _, h, w = masks.shape77 device = masks.device78 xs = torch.arange(w, device=device, dtype=torch.int32)79 ys = torch.arange(h, device=device, dtype=torch.int32)80 grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing="xy")81 grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w)82 grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w)83 min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1)84 max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1)85 min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1)86 max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1)87 bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1)88 89 return bbox_coords90 91 92def _load_img_as_tensor(img_path, image_size):93 img_pil = Image.open(img_path)94 img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size)))95 if img_np.dtype == np.uint8: # np.uint8 is expected for JPEG images96 img_np = img_np / 255.097 else:98 raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}")99 img = torch.from_numpy(img_np).permute(2, 0, 1)100 video_width, video_height = img_pil.size # the original video size101 return img, video_height, video_width102 103 104class AsyncVideoFrameLoader:105 """106 A list of video frames to be load asynchronously without blocking session start.107 """108 109 def __init__(110 self,111 img_paths,112 image_size,113 offload_video_to_cpu,114 img_mean,115 img_std,116 compute_device,117 ):118 self.img_paths = img_paths119 self.image_size = image_size120 self.offload_video_to_cpu = offload_video_to_cpu121 self.img_mean = img_mean122 self.img_std = img_std123 # items in `self.images` will be loaded asynchronously124 self.images = [None] * len(img_paths)125 # catch and raise any exceptions in the async loading thread126 self.exception = None127 # video_height and video_width be filled when loading the first image128 self.video_height = None129 self.video_width = None130 self.compute_device = compute_device131 132 # load the first frame to fill video_height and video_width and also133 # to cache it (since it's most likely where the user will click)134 self.__getitem__(0)135 136 # load the rest of frames asynchronously without blocking the session start137 def _load_frames():138 try:139 for n in tqdm(range(len(self.images)), desc="frame loading (JPEG)"):140 self.__getitem__(n)141 except Exception as e:142 self.exception = e143 144 self.thread = Thread(target=_load_frames, daemon=True)145 self.thread.start()146 147 def __getitem__(self, index):148 if self.exception is not None:149 raise RuntimeError("Failure in frame loading thread") from self.exception150 151 img = self.images[index]152 if img is not None:153 return img154 155 img, video_height, video_width = _load_img_as_tensor(156 self.img_paths[index], self.image_size157 )158 self.video_height = video_height159 self.video_width = video_width160 # normalize by mean and std161 img -= self.img_mean162 img /= self.img_std163 if not self.offload_video_to_cpu:164 img = img.to(self.compute_device, non_blocking=True)165 self.images[index] = img166 return img167 168 def __len__(self):169 return len(self.images)170 171 172def load_video_frames(173 video_path,174 image_size,175 offload_video_to_cpu,176 img_mean=(0.485, 0.456, 0.406),177 img_std=(0.229, 0.224, 0.225),178 async_loading_frames=False,179 compute_device=torch.device("cuda"),180):181 """182 Load the video frames from video_path. The frames are resized to image_size as in183 the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo.184 """185 is_bytes = isinstance(video_path, bytes)186 is_str = isinstance(video_path, str)187 is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [".mp4", ".MP4"]188 if is_bytes or is_mp4_path:189 return load_video_frames_from_video_file(190 video_path=video_path,191 image_size=image_size,192 offload_video_to_cpu=offload_video_to_cpu,193 img_mean=img_mean,194 img_std=img_std,195 compute_device=compute_device,196 )197 elif is_str and os.path.isdir(video_path):198 return load_video_frames_from_jpg_images(199 video_path=video_path,200 image_size=image_size,201 offload_video_to_cpu=offload_video_to_cpu,202 img_mean=img_mean,203 img_std=img_std,204 async_loading_frames=async_loading_frames,205 compute_device=compute_device,206 )207 else:208 raise NotImplementedError(209 "Only MP4 video and JPEG folder are supported at this moment"210 )211 212 213def load_video_frames_from_jpg_images(214 video_path,215 image_size,216 offload_video_to_cpu,217 img_mean=(0.485, 0.456, 0.406),218 img_std=(0.229, 0.224, 0.225),219 async_loading_frames=False,220 compute_device=torch.device("cuda"),221):222 """223 Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format).224 225 The frames are resized to image_size x image_size and are loaded to GPU if226 `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`.227 228 You can load a frame asynchronously by setting `async_loading_frames` to `True`.229 """230 if isinstance(video_path, str) and os.path.isdir(video_path):231 jpg_folder = video_path232 else:233 raise NotImplementedError(234 "Only JPEG frames are supported at this moment. For video files, you may use "235 "ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \n"236 "```\n"237 "ffmpeg -i <your_video>.mp4 -q:v 2 -start_number 0 <output_dir>/'%05d.jpg'\n"238 "```\n"239 "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks "240 "ffmpeg to start the JPEG file from 00000.jpg."241 )242 243 frame_names = [244 p245 for p in os.listdir(jpg_folder)246 if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]247 ]248 frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))249 num_frames = len(frame_names)250 if num_frames == 0:251 raise RuntimeError(f"no images found in {jpg_folder}")252 img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names]253 img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]254 img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]255 256 if async_loading_frames:257 lazy_images = AsyncVideoFrameLoader(258 img_paths,259 image_size,260 offload_video_to_cpu,261 img_mean,262 img_std,263 compute_device,264 )265 return lazy_images, lazy_images.video_height, lazy_images.video_width266 267 images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)268 for n, img_path in enumerate(tqdm(img_paths, desc="frame loading (JPEG)")):269 images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)270 if not offload_video_to_cpu:271 images = images.to(compute_device)272 img_mean = img_mean.to(compute_device)273 img_std = img_std.to(compute_device)274 # normalize by mean and std275 images -= img_mean276 images /= img_std277 return images, video_height, video_width278 279 280def load_video_frames_from_video_file(281 video_path,282 image_size,283 offload_video_to_cpu,284 img_mean=(0.485, 0.456, 0.406),285 img_std=(0.229, 0.224, 0.225),286 compute_device=torch.device("cuda"),287):288 """Load the video frames from a video file."""289 import decord290 291 img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]292 img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]293 # Get the original video height and width294 decord.bridge.set_bridge("torch")295 video_height, video_width, _ = decord.VideoReader(video_path).next().shape296 # Iterate over all frames in the video297 images = []298 for frame in decord.VideoReader(video_path, width=image_size, height=image_size):299 images.append(frame.permute(2, 0, 1))300 301 images = torch.stack(images, dim=0).float() / 255.0302 if not offload_video_to_cpu:303 images = images.to(compute_device)304 img_mean = img_mean.to(compute_device)305 img_std = img_std.to(compute_device)306 # normalize by mean and std307 images -= img_mean308 images /= img_std309 return images, video_height, video_width310 311 312def fill_holes_in_mask_scores(mask, max_area):313 """314 A post processor to fill small holes in mask scores with area under `max_area`.315 """316 # Holes are those connected components in background with area <= self.max_area317 # (background regions are those with mask scores <= 0)318 assert max_area > 0, "max_area must be positive"319 320 input_mask = mask321 try:322 labels, areas = get_connected_components(mask <= 0)323 is_hole = (labels > 0) & (areas <= max_area)324 # We fill holes with a small positive mask score (0.1) to change them to foreground.325 mask = torch.where(is_hole, 0.1, mask)326 except Exception as e:327 # Skip the post-processing step on removing small holes if the CUDA kernel fails328 warnings.warn(329 f"{e}\n\nSkipping the post-processing step due to the error above. You can "330 "still use SAM 2 and it's OK to ignore the error above, although some post-processing "331 "functionality may be limited (which doesn't affect the results in most cases; see "332 "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).",333 category=UserWarning,334 stacklevel=2,335 )336 mask = input_mask337 338 return mask339 340 341def concat_points(old_point_inputs, new_points, new_labels):342 """Add new points and labels to previous point inputs (add at the end)."""343 if old_point_inputs is None:344 points, labels = new_points, new_labels345 else:346 points = torch.cat([old_point_inputs["point_coords"], new_points], dim=1)347 labels = torch.cat([old_point_inputs["point_labels"], new_labels], dim=1)348 349 return {"point_coords": points, "point_labels": labels}350