| | """ |
| | This space is modified from the HF space of fffiloni/SAM2-Video-Predictor |
| | """ |
| |
|
| | import subprocess |
| | import re |
| | from typing import List, Tuple, Optional |
| | import pickle |
| |
|
| | |
| | command = ["python", "setup.py", "build_ext", "--inplace"] |
| |
|
| | |
| | result = subprocess.run(command, capture_output=True, text=True) |
| |
|
| | |
| | print("Output:\n", result.stdout) |
| | print("Errors:\n", result.stderr) |
| |
|
| | |
| | if result.returncode == 0: |
| | print("Command executed successfully.") |
| | else: |
| | print("Command failed with return code:", result.returncode) |
| |
|
| | import gradio as gr |
| | from datetime import datetime |
| | import os |
| | os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1" |
| | import torch |
| | import numpy as np |
| | import cv2 |
| | import matplotlib.pyplot as plt |
| | from PIL import Image, ImageFilter |
| | from sam2.build_sam import build_sam2_video_predictor |
| |
|
| | from moviepy.editor import ImageSequenceClip |
| |
|
| | def get_video_properties(video_path): |
| | |
| | cap = cv2.VideoCapture(video_path) |
| | |
| | if not cap.isOpened(): |
| | print("Error: Could not open video.") |
| | return None |
| | |
| | |
| | fps = cap.get(cv2.CAP_PROP_FPS) |
| | |
| | |
| | width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| | height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| |
|
| | |
| | cap.release() |
| |
|
| | return fps, width, height |
| |
|
| | def clear_points(image): |
| | |
| | return [ |
| | image, |
| | gr.State([]), |
| | gr.State([]), |
| | image, |
| | |
| | ] |
| |
|
| | def preprocess_video_in(video_path): |
| |
|
| | |
| | unique_id = datetime.now().strftime('%Y%m%d%H%M%S') |
| | |
| | |
| | extracted_frames_output_dir = f'frames_{unique_id}' |
| | |
| | |
| | os.makedirs(extracted_frames_output_dir, exist_ok=True) |
| |
|
| | |
| | |
| | cap = cv2.VideoCapture(video_path) |
| | |
| | if not cap.isOpened(): |
| | print("Error: Could not open video.") |
| | return None |
| |
|
| | |
| | fps = cap.get(cv2.CAP_PROP_FPS) |
| | |
| | |
| | max_frames = int(fps * 10) |
| | |
| | frame_number = 0 |
| | first_frame = None |
| | |
| | while True: |
| | ret, frame = cap.read() |
| | """ |
| | 不再将视频裁剪至10s |
| | if not ret or frame_number >= max_frames: |
| | break |
| | """ |
| | if not ret: |
| | break |
| |
|
| | |
| | |
| | frame_filename = os.path.join(extracted_frames_output_dir, f'{frame_number:05d}.jpg') |
| | |
| | |
| | cv2.imwrite(frame_filename, frame) |
| | |
| | |
| | if frame_number == 0: |
| | first_frame = frame_filename |
| | |
| | frame_number += 1 |
| | |
| | |
| | cap.release() |
| | |
| | |
| | scanned_frames = [ |
| | p for p in os.listdir(extracted_frames_output_dir) |
| | if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] |
| | ] |
| | scanned_frames.sort(key=lambda p: int(os.path.splitext(p)[0])) |
| | |
| | |
| | return [ |
| | first_frame, |
| | gr.State([]), |
| | gr.State([]), |
| | first_frame, |
| | first_frame, |
| | extracted_frames_output_dir, |
| | scanned_frames, |
| | None, |
| | None, |
| | gr.update(open=False) |
| | ] |
| |
|
| | def get_point(point_type, tracking_points, trackings_input_label, input_first_frame_image, evt: gr.SelectData): |
| | print(f"You selected {evt.value} at {evt.index} from {evt.target}") |
| |
|
| | tracking_points.value.append(evt.index) |
| | print(f"TRACKING POINT: {tracking_points.value}") |
| |
|
| | if point_type == "include": |
| | trackings_input_label.value.append(1) |
| | elif point_type == "exclude": |
| | trackings_input_label.value.append(0) |
| | print(f"TRACKING INPUT LABEL: {trackings_input_label.value}") |
| | |
| | |
| | transparent_background = Image.open(input_first_frame_image).convert('RGBA') |
| | w, h = transparent_background.size |
| | |
| | |
| | fraction = 0.02 |
| | radius = int(fraction * min(w, h)) |
| | |
| | |
| | transparent_layer = np.zeros((h, w, 4), dtype=np.uint8) |
| | |
| | for index, track in enumerate(tracking_points.value): |
| | if trackings_input_label.value[index] == 1: |
| | cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1) |
| | else: |
| | cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1) |
| |
|
| | |
| | transparent_layer = Image.fromarray(transparent_layer, 'RGBA') |
| | selected_point_map = Image.alpha_composite(transparent_background, transparent_layer) |
| | |
| | return tracking_points, trackings_input_label, selected_point_map |
| | |
| | |
| | torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__() |
| |
|
| | if torch.cuda.get_device_properties(0).major >= 8: |
| | |
| | torch.backends.cuda.matmul.allow_tf32 = True |
| | torch.backends.cudnn.allow_tf32 = True |
| | |
| | def show_mask(mask, ax, obj_id=None, random_color=False): |
| | if random_color: |
| | color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) |
| | else: |
| | cmap = plt.get_cmap("tab10") |
| | cmap_idx = 0 if obj_id is None else obj_id |
| | color = np.array([*cmap(cmap_idx)[:3], 0.6]) |
| | h, w = mask.shape[-2:] |
| | mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) |
| | ax.imshow(mask_image) |
| |
|
| | def show_white_mask(mask, ax): |
| | |
| | mask = (mask > 0).astype(float) |
| | h, w = mask.shape[-2:] |
| | |
| | |
| | alpha = 1.0 |
| | color = np.array([1, 1, 1, alpha]) |
| | mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1) |
| | |
| | |
| | ax.imshow(np.zeros((h, w, 3), dtype=float)) |
| | ax.imshow(mask_image) |
| |
|
| | def show_points(coords, labels, ax, marker_size=200): |
| | pos_points = coords[labels==1] |
| | neg_points = coords[labels==0] |
| | ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) |
| | ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) |
| |
|
| | def show_box(box, ax): |
| | x0, y0 = box[0], box[1] |
| | w, h = box[2] - box[0], box[3] - box[1] |
| | ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2)) |
| |
|
| |
|
| | def load_model(checkpoint): |
| | |
| | if checkpoint == "tiny": |
| | sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt" |
| | model_cfg = "sam2_hiera_t.yaml" |
| | return [sam2_checkpoint, model_cfg] |
| | elif checkpoint == "samll": |
| | sam2_checkpoint = "./checkpoints/sam2_hiera_small.pt" |
| | model_cfg = "sam2_hiera_s.yaml" |
| | return [sam2_checkpoint, model_cfg] |
| | elif checkpoint == "base-plus": |
| | sam2_checkpoint = "./checkpoints/sam2_hiera_base_plus.pt" |
| | model_cfg = "sam2_hiera_b+.yaml" |
| | return [sam2_checkpoint, model_cfg] |
| | elif checkpoint == "large": |
| | sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt" |
| | model_cfg = "sam2_hiera_l.yaml" |
| | return [sam2_checkpoint, model_cfg] |
| |
|
| | |
| | |
| | def get_mask_sam_process( |
| | stored_inference_state, |
| | input_first_frame_image, |
| | checkpoint, |
| | tracking_points, |
| | trackings_input_label, |
| | video_frames_dir, |
| | scanned_frames, |
| | working_frame: str = None, |
| | available_frames_to_check: List[str] = [], |
| | |
| | ): |
| | |
| | |
| | print(f"USER CHOSEN CHECKPOINT: {checkpoint}") |
| | sam2_checkpoint, model_cfg = load_model(checkpoint) |
| | print("MODEL LOADED") |
| |
|
| | |
| | predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint) |
| | print("PREDICTOR READY") |
| |
|
| | |
| | |
| | video_dir = video_frames_dir |
| | |
| | |
| | frame_names = scanned_frames |
| |
|
| | |
| | if stored_inference_state is None: |
| | |
| | inference_state = predictor.init_state(video_path=video_dir) |
| | print("NEW INFERENCE_STATE INITIATED") |
| | else: |
| | inference_state = stored_inference_state |
| |
|
| | |
| | |
| |
|
| | |
| | |
| | |
| | if working_frame is None: |
| | ann_frame_idx = 0 |
| | working_frame = "frame_0.jpg" |
| | else: |
| | |
| | match = re.search(r'frame_(\d+)', working_frame) |
| | if match: |
| | |
| | frame_number = int(match.group(1)) |
| | ann_frame_idx = frame_number |
| | |
| | print(f"NEW_WORKING_FRAME PATH: {working_frame}") |
| | |
| | ann_obj_id = 1 |
| | |
| | |
| | points = np.array(tracking_points.value, dtype=np.float32) |
| | |
| | labels = np.array(trackings_input_label.value, np.int32) |
| | _, out_obj_ids, out_mask_logits = predictor.add_new_points( |
| | inference_state=inference_state, |
| | frame_idx=ann_frame_idx, |
| | obj_id=ann_obj_id, |
| | points=points, |
| | labels=labels, |
| | ) |
| |
|
| | |
| | plt.figure(figsize=(12, 8)) |
| | plt.title(f"frame {ann_frame_idx}") |
| | plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx]))) |
| | show_points(points, labels, plt.gca()) |
| | show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0]) |
| | |
| | |
| | first_frame_output_filename = "output_first_frame.jpg" |
| | plt.savefig(first_frame_output_filename, format='jpg') |
| | plt.close() |
| | torch.cuda.empty_cache() |
| |
|
| | |
| | if working_frame not in available_frames_to_check: |
| | available_frames_to_check.append(working_frame) |
| | print(available_frames_to_check) |
| | |
| | return gr.update(visible=True), "output_first_frame.jpg", frame_names, predictor, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=True) |
| |
|
| | def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame, progress=gr.Progress(track_tqdm=True)): |
| | |
| | sam2_checkpoint, model_cfg = load_model(checkpoint) |
| | predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint) |
| | |
| | inference_state = stored_inference_state |
| | frame_names = stored_frame_names |
| | video_dir = video_frames_dir |
| | |
| | |
| | frames_output_dir = "frames_output_images" |
| | os.makedirs(frames_output_dir, exist_ok=True) |
| |
|
| | mask_frames_output_dir = "mask_frames_output_images" |
| | os.makedirs(mask_frames_output_dir, exist_ok=True) |
| | |
| | |
| | jpeg_images = [] |
| | masks_images = [] |
| |
|
| | |
| | |
| | global video_segments |
| | video_segments = {} |
| | 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) |
| | } |
| | |
| | |
| | if vis_frame_type == "check": |
| | vis_frame_stride = 15 |
| | elif vis_frame_type == "render": |
| | vis_frame_stride = 1 |
| | |
| | plt.close("all") |
| | for out_frame_idx in range(0, len(frame_names), vis_frame_stride): |
| | plt.figure(figsize=(6, 4)) |
| | plt.title(f"frame {out_frame_idx}") |
| | plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx]))) |
| | for out_obj_id, out_mask in video_segments[out_frame_idx].items(): |
| | show_mask(out_mask, plt.gca(), obj_id=out_obj_id) |
| |
|
| | |
| | output_filename = os.path.join(frames_output_dir, f"frame_{out_frame_idx}.jpg") |
| | plt.savefig(output_filename, format='jpg') |
| | |
| | |
| | plt.close() |
| |
|
| | |
| | jpeg_images.append(output_filename) |
| |
|
| | if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check: |
| | available_frames_to_check.append(f"frame_{out_frame_idx}.jpg") |
| |
|
| | |
| | |
| | fig, ax = plt.subplots(figsize=(6, 4)) |
| | ax.axis("off") |
| | |
| | |
| | for out_mask in video_segments[out_frame_idx].values(): |
| | show_white_mask(out_mask, ax) |
| | |
| | |
| | mask_filename = os.path.join(mask_frames_output_dir, f"mask_{out_frame_idx}.jpg") |
| | fig.savefig(mask_filename, format='jpg', bbox_inches="tight", pad_inches=0) |
| | plt.close(fig) |
| | |
| | |
| | masks_images.append(mask_filename) |
| |
|
| | |
| |
|
| | torch.cuda.empty_cache() |
| | print(f"JPEG_IMAGES: {jpeg_images}") |
| |
|
| | if vis_frame_type == "check": |
| | return gr.update(value=jpeg_images), gr.update(value=None), gr.update(choices=available_frames_to_check, value=working_frame, visible=True), available_frames_to_check, gr.update(visible=True), None |
| | elif vis_frame_type == "render": |
| | |
| | original_fps, o_width, o_height = get_video_properties(video_in) |
| | fps = original_fps |
| | total_frames = len(jpeg_images) |
| | clip = ImageSequenceClip(jpeg_images, fps=fps) |
| | |
| | final_vid_output_path = "output_video.mp4" |
| | |
| | |
| | clip.write_videofile( |
| | final_vid_output_path, |
| | codec='libx264' |
| | ) |
| |
|
| | print("MAKING MASK VIDEO ...") |
| |
|
| | |
| | mask_video_filename = "final_masks_video.mp4" |
| | |
| | |
| | fourcc = cv2.VideoWriter_fourcc(*"mp4v") |
| | video_writer = cv2.VideoWriter(mask_video_filename, fourcc, fps, (o_width, o_height)) |
| | |
| | |
| | for mask_path in masks_images: |
| | |
| | frame = cv2.imread(mask_path) |
| | |
| | |
| | frame_resized = cv2.resize(frame, (o_width, o_height), interpolation=cv2.INTER_NEAREST) |
| | |
| | |
| | video_writer.write(frame_resized) |
| | |
| | |
| | video_writer.release() |
| | print(f"Mask Video saved at {mask_video_filename}") |
| |
|
| | |
| | return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True), mask_video_filename |
| |
|
| | def update_ui(vis_frame_type): |
| | if vis_frame_type == "check": |
| | return gr.update(visible=True), gr.update(visible=False) |
| | elif vis_frame_type == "render": |
| | return gr.update(visible=False), gr.update(visible=True) |
| |
|
| | def switch_working_frame(working_frame, scanned_frames, video_frames_dir): |
| | new_working_frame = None |
| | if working_frame == None: |
| | new_working_frame = os.path.join(video_frames_dir, scanned_frames[0]) |
| | |
| | else: |
| | |
| | match = re.search(r'frame_(\d+)', working_frame) |
| | if match: |
| | |
| | frame_number = int(match.group(1)) |
| | ann_frame_idx = frame_number |
| | new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx]) |
| | return gr.State([]), gr.State([]), new_working_frame, new_working_frame |
| |
|
| | def reset_propagation(first_frame_path, predictor, stored_inference_state): |
| | |
| | predictor.reset_state(stored_inference_state) |
| | |
| | return first_frame_path, gr.State([]), gr.State([]), gr.update(value=None, visible=False), stored_inference_state, None, ["frame_0.jpg"], first_frame_path, "frame_0.jpg", gr.update(visible=False) |
| |
|
| | def export_masks(): |
| | |
| | global video_segments |
| | if not video_segments: |
| | raise gr.Error("No mask data available. Please run propagation first!") |
| | |
| | |
| | filename = "mask_data.pkl" |
| | with open(filename, 'wb') as f: |
| | pickle.dump(video_segments, f) |
| | |
| | return filename |
| |
|
| | css=""" |
| | div#component-18, div#component-25, div#component-35, div#component-41{ |
| | align-items: stretch!important; |
| | } |
| | """ |
| |
|
| | with gr.Blocks(css=css) as demo: |
| | first_frame_path = gr.State() |
| | tracking_points = gr.State([]) |
| | trackings_input_label = gr.State([]) |
| | video_frames_dir = gr.State() |
| | scanned_frames = gr.State() |
| | loaded_predictor = gr.State() |
| | stored_inference_state = gr.State() |
| | stored_frame_names = gr.State() |
| | available_frames_to_check = gr.State([]) |
| | with gr.Column(): |
| | gr.Markdown("# SAM2 Video Predictor") |
| | gr.Markdown("This is a simple demo for video segmentation with SAM2.") |
| | |
| | with gr.Row(): |
| | |
| | with gr.Column(): |
| | with gr.Group(): |
| | with gr.Group(): |
| | with gr.Row(): |
| | point_type = gr.Radio(label="point type", choices=["include", "exclude"], value="include", scale=2) |
| | clear_points_btn = gr.Button("Clear Points", scale=1) |
| | |
| | input_first_frame_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False) |
| | |
| | points_map = gr.Image( |
| | label="Point n Click map", |
| | type="filepath", |
| | interactive=False |
| | ) |
| | |
| | with gr.Group(): |
| | with gr.Row(): |
| | checkpoint = gr.Dropdown(label="Checkpoint", choices=["tiny", "small", "base-plus", "large"], value="tiny") |
| | submit_btn = gr.Button("Get Mask", size="lg") |
| |
|
| | with gr.Accordion("Your video IN", open=True) as video_in_drawer: |
| | video_in = gr.Video(label="Video IN", format="mp4") |
| |
|
| | with gr.Group(): |
| | with gr.Group(): |
| | with gr.Row(): |
| | working_frame = gr.Dropdown(label="working frame ID", choices=[""], value=None, visible=False, allow_custom_value=False, interactive=True) |
| | change_current = gr.Button("change current", visible=False) |
| | output_result = gr.Image(label="current working mask ref") |
| | with gr.Group(): |
| | with gr.Row(): |
| | vis_frame_type = gr.Radio(label="Propagation level", choices=["check", "render"], value="check", scale=2) |
| | propagate_btn = gr.Button("Propagate", scale=1) |
| | reset_prpgt_brn = gr.Button("Reset", visible=False) |
| | |
| | gr.HTML(""" |
| | |
| | <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true"> |
| | <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" /> |
| | </a> to skip queue and avoid OOM errors from heavy public load |
| | """) |
| | |
| | with gr.Column(): |
| |
|
| | gr.Markdown(""" |
| | This space is modified from the HF space of fffiloni/SAM2-Video-Predictor |
| | |
| | Instructions: |
| | 1. Upload your video in .mp4 format |
| | 2. Add points (include & exclude) to the first frame |
| | 3. Get mask to have a preview of the mask |
| | 4. Render propagate to output the final mask |
| | 5. A button will appear at the bottom of the webpage, click it to download the mask in pickle format |
| | """) |
| | |
| | output_propagated = gr.Gallery(label="Propagated Mask samples gallery", columns=4, visible=False) |
| | output_video = gr.Video(visible=False) |
| | mask_final_output = gr.Video(label="Mask Video") |
| | |
| | |
| | export_btn = gr.Button("Export Mask Data") |
| | mask_download = gr.File(label="Download Mask Data", visible=False) |
| | |
| | |
| |
|
| | |
| | video_in.upload( |
| | fn = preprocess_video_in, |
| | inputs = [video_in], |
| | outputs = [ |
| | first_frame_path, |
| | tracking_points, |
| | trackings_input_label, |
| | input_first_frame_image, |
| | points_map, |
| | video_frames_dir, |
| | scanned_frames, |
| | stored_inference_state, |
| | stored_frame_names, |
| | video_in_drawer, |
| | ], |
| | queue = False |
| | ) |
| |
|
| | |
| | |
| | points_map.select( |
| | fn = get_point, |
| | inputs = [ |
| | point_type, |
| | tracking_points, |
| | trackings_input_label, |
| | input_first_frame_image, |
| | ], |
| | outputs = [ |
| | tracking_points, |
| | trackings_input_label, |
| | points_map, |
| | ], |
| | queue = False |
| | ) |
| |
|
| | |
| | clear_points_btn.click( |
| | fn = clear_points, |
| | inputs = input_first_frame_image, |
| | outputs = [ |
| | first_frame_path, |
| | tracking_points, |
| | trackings_input_label, |
| | points_map, |
| | |
| | ], |
| | queue=False |
| | ) |
| |
|
| | |
| | change_current.click( |
| | fn = switch_working_frame, |
| | inputs = [working_frame, scanned_frames, video_frames_dir], |
| | outputs = [tracking_points, trackings_input_label, input_first_frame_image, points_map], |
| | queue=False |
| | ) |
| | |
| |
|
| | submit_btn.click( |
| | fn = get_mask_sam_process, |
| | inputs = [ |
| | stored_inference_state, |
| | input_first_frame_image, |
| | checkpoint, |
| | tracking_points, |
| | trackings_input_label, |
| | video_frames_dir, |
| | scanned_frames, |
| | working_frame, |
| | available_frames_to_check, |
| | ], |
| | outputs = [ |
| | change_current, |
| | output_result, |
| | stored_frame_names, |
| | loaded_predictor, |
| | stored_inference_state, |
| | working_frame, |
| | ], |
| | queue=False |
| | ) |
| |
|
| | reset_prpgt_brn.click( |
| | fn = reset_propagation, |
| | inputs = [first_frame_path, loaded_predictor, stored_inference_state], |
| | outputs = [points_map, tracking_points, trackings_input_label, output_propagated, stored_inference_state, output_result, available_frames_to_check, input_first_frame_image, working_frame, reset_prpgt_brn], |
| | queue=False |
| | ) |
| |
|
| | propagate_btn.click( |
| | fn = update_ui, |
| | inputs = [vis_frame_type], |
| | outputs = [output_propagated, output_video], |
| | queue=False |
| | ).then( |
| | fn = propagate_to_all, |
| | inputs = [video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame], |
| | outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn, mask_final_output] |
| | ) |
| |
|
| | |
| | export_btn.click( |
| | fn=export_masks, |
| | outputs=[mask_download] |
| | ).then( |
| | fn=lambda: gr.update(visible=True), |
| | outputs=[mask_download] |
| | ) |
| |
|
| | demo.launch(show_api=False, show_error=True) |