Buckets:
| import torch | |
| import torchvision | |
| print("PyTorch version:", torch.__version__) | |
| print("Torchvision version:", torchvision.__version__) | |
| print("CUDA is available:", torch.cuda.is_available()) | |
| import os | |
| from pathlib import Path | |
| import cv2 | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from PIL import Image | |
| from transformers import Sam3VideoModel, Sam3VideoProcessor | |
| from accelerate import Accelerator | |
| from transformers.video_utils import load_video | |
| device = Accelerator().device | |
| model = Sam3VideoModel.from_pretrained("facebook/sam3").to(device, dtype=torch.bfloat16) | |
| processor = Sam3VideoProcessor.from_pretrained("facebook/sam3") | |
| def save_frame_detections_fast(image_source, out_dict, output_path): | |
| """ | |
| Blazing fast bounding box rendering and saving using OpenCV. | |
| Args: | |
| image_source: Path to the image file (str) OR a pre-loaded BGR numpy array. | |
| out_dict: Dict containing 'out_obj_ids', 'out_probs', 'out_boxes_xywh', and 'frame_stats'. | |
| output_path: Destination path to save the annotated image. | |
| """ | |
| # 1. Load image efficiently (skip decoding if already a numpy array) | |
| if isinstance(image_source, str): | |
| img = cv2.imread(image_source) | |
| else: | |
| img = image_source[:, :, ::-1].copy() # Copy to avoid mutating original video frame buffer | |
| H, W, _ = img.shape | |
| # 2. Extract arrays from dict | |
| ids = out_dict['object_ids'].cpu().numpy() | |
| probs = out_dict['scores'].cpu().numpy() | |
| boxes = out_dict['boxes'].cpu().numpy() | |
| # stats = out_dict.get('frame_stats', {}) | |
| # 3. Vectorized normalization conversion (XYWH normalized -> Absolute XYXY) | |
| # Assumes XYWH is [top_left_x, top_left_y, width, height] | |
| # x1 = (boxes[:, 0] * W).astype(np.int32) | |
| # y1 = (boxes[:, 1] * H).astype(np.int32) | |
| # x2 = ((boxes[:, 0] + boxes[:, 2]) * W).astype(np.int32) | |
| # y2 = ((boxes[:, 1] + boxes[:, 3]) * H).astype(np.int32) | |
| # Note: If your model outputs CENTER-based XYWH [cx, cy, w, h], use this instead: | |
| # x1 = ((boxes[:, 0] - boxes[:, 2] / 2) * W).astype(np.int32) | |
| # y1 = ((boxes[:, 1] - boxes[:, 3] / 2) * H).astype(np.int32) | |
| # x2 = ((boxes[:, 0] + boxes[:, 2] / 2) * W).astype(np.int32) | |
| # y2 = ((boxes[:, 1] + boxes[:, 3] / 2) * H).astype(np.int32) | |
| # 4. Draw Detections | |
| for i in range(len(ids)): | |
| obj_id = int(ids[i]) | |
| score = probs[i] | |
| x1,y1,x2,y2 = boxes[i].astype(np.int32) | |
| # Fast, deterministic pseudo-random color generation based on Obj ID (BGR format) | |
| color = (0, 0, 255) | |
| # Draw Bounding Box | |
| cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness=2) | |
| # Render ID & Score Text | |
| label = f"ID: {obj_id} ({score:.2f})" | |
| # Place text slightly above the box, or inside if too close to the top edge | |
| text_y = max(y1 - 7, 20) | |
| cv2.putText( | |
| img, label, (x1, text_y), | |
| fontFace=cv2.FONT_HERSHEY_SIMPLEX, | |
| fontScale=0.5, color=color, thickness=2, lineType=cv2.LINE_AA | |
| ) | |
| # 5. Draw Global Frame Stats (Top Left Corner Overlay) | |
| # tracked = stats.get('num_obj_tracked', 0) | |
| # dropped = stats.get('num_obj_dropped', 0) | |
| # stats_text = f"Tracked: {tracked} | Dropped: {dropped}" | |
| # Background banner for stats readability | |
| # cv2.rectangle(img, (0, 0), (320, 35), (0, 0, 0), -1) | |
| # cv2.putText( | |
| # img, stats_text, (10, 23), | |
| # fontFace=cv2.FONT_HERSHEY_SIMPLEX, | |
| # fontScale=0.6, color=(255, 255, 255), thickness=2, lineType=cv2.LINE_AA | |
| # ) | |
| # 6. Save directly to disk | |
| os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) | |
| cv2.imwrite(output_path, img) | |
| def main(): | |
| # font size for axes titles | |
| plt.rcParams["axes.titlesize"] = 12 | |
| plt.rcParams["figure.titlesize"] = 12 | |
| # "video_path" needs to be either a JPEG folder or a MP4 video file | |
| #video_path = f"{sam3_root}/assets/videos/0001" | |
| video_path = str(Path(__file__).parent / "input.mp4") | |
| output_path = Path(__file__).parent / "out_predictor" | |
| video_name = Path(video_path).stem | |
| video_output_dir = output_path / video_name | |
| video_output_dir.mkdir(parents=True, exist_ok=True) | |
| video_frames_for_vis, _ = load_video(video_path,fps=1) | |
| pil_image_list = [Image.fromarray(img) for img in video_frames_for_vis] | |
| inference_session = processor.init_video_session( | |
| video=pil_image_list, | |
| inference_device=device, | |
| processing_device="cpu", | |
| video_storage_device="cpu", | |
| dtype=torch.bfloat16, | |
| ) | |
| keywords = ["vehicle"] | |
| for prompt in keywords: | |
| # create output folder for the current prompt | |
| prompt_output_dir = video_output_dir / prompt | |
| prompt_output_dir.mkdir(exist_ok=True) | |
| inference_session = processor.add_text_prompt( | |
| inference_session=inference_session, | |
| text=prompt, | |
| ) | |
| for model_outputs in model.propagate_in_video_iterator( | |
| inference_session=inference_session, max_frame_num_to_track=len(pil_image_list) | |
| ): | |
| processed_outputs = processor.postprocess_outputs(inference_session, model_outputs) | |
| frame_idx = model_outputs.frame_idx | |
| if len(processed_outputs['scores']) > 0: | |
| save_frame_detections_fast(np.array(pil_image_list[frame_idx]), processed_outputs, str(prompt_output_dir / f"frame_{frame_idx}.png")) | |
| if __name__ == '__main__': | |
| main() |
Xet Storage Details
- Size:
- 5.46 kB
- Xet hash:
- 03eed563d6f10faf0b93e297013c942772ae0f40e29d56f80ad0c00363897a85
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.