Update app.py

app.py

@@ -1,178 +1,428 @@
-# The full rewritten version of the provided code with progress bar, error fixes, and proper Gradio integration
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
 
-import os
 import copy
-import tempfile
+import os
 from datetime import datetime
-import gc
+import tempfile
 
 import cv2
-import numpy as np
-from PIL import Image
 import matplotlib.pyplot as plt
-import torch
+import numpy as np
 import gradio as gr
+import torch
 from moviepy.editor import ImageSequenceClip
-
+from PIL import Image
 from sam2.build_sam import build_sam2_video_predictor
 
-# Remove CUDA-related env var to force CPU-only mode
-os.environ.pop("TORCH_CUDNN_SDPA_ENABLED", None)
+# Remove CUDA environment variables
+if 'TORCH_CUDNN_SDPA_ENABLED' in os.environ:
+    del os.environ["TORCH_CUDNN_SDPA_ENABLED"]
+
+# UI Description
+title = "<center><strong><font size='8'>EdgeTAM CPU</font></strong> <a href='https://github.com/facebookresearch/EdgeTAM'><font size='6'>[GitHub]</font></a></center>"
+
+description_p = """# Instructions
+                <ol>
+                <li>Upload one video or click one example video</li>
+                <li>Click 'include' point type, select the object to segment and track</li>
+                <li>Click 'exclude' point type (optional), select the area to avoid segmenting</li>
+                <li>Click the 'Track' button to obtain the masked video</li>
+                </ol>
+              """
+
+# Example videos
+examples = [
+    ["examples/01_dog.mp4"],
+    ["examples/02_cups.mp4"],
+    ["examples/03_blocks.mp4"],
+    ["examples/04_coffee.mp4"],
+    ["examples/05_default_juggle.mp4"],
+]
+
+OBJ_ID = 0
 
-# Config
+# Initialize model on CPU
 sam2_checkpoint = "checkpoints/edgetam.pt"
 model_cfg = "edgetam.yaml"
-examples = [[f"examples/{vid}"] for vid in ["01_dog.mp4", "02_cups.mp4", "03_blocks.mp4", "04_coffee.mp4", "05_default_juggle.mp4"]]
-OBJ_ID = 0
 
-# Model loader
-if os.path.exists(sam2_checkpoint) and os.path.exists(model_cfg):
-    try:
-        predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
-    except Exception as e:
-        print("Error loading predictor:", e)
-        predictor = None
-else:
-    print("Model files missing.")
+def check_file_exists(filepath):
+    exists = os.path.exists(filepath)
+    if not exists:
+        print(f"WARNING: File not found: {filepath}")
+    return exists
+
+# Verify model files
+model_files_exist = check_file_exists(sam2_checkpoint) and check_file_exists(model_cfg)
+try:
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
+    print("Predictor loaded on CPU")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    import traceback
+    traceback.print_exc()
     predictor = None
 
-def get_fps(video_path):
+# Utility Functions
+def get_video_fps(video_path):
     cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened(): return 30.0
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return 30.0
     fps = cap.get(cv2.CAP_PROP_FPS)
     cap.release()
     return fps
 
-def reset(session):
-    if session["inference_state"]:
-        predictor.reset_state(session["inference_state"])
-    session.update({"input_points": [], "input_labels": [], "first_frame": None, "all_frames": None, "inference_state": None})
-    return None, gr.update(open=True), None, None, gr.update(value=None, visible=False), session
+def reset(session_state):
+    session_state["input_points"] = []
+    session_state["input_labels"] = []
+    if session_state["inference_state"] is not None:
+        predictor.reset_state(session_state["inference_state"])
+    session_state["first_frame"] = None
+    session_state["all_frames"] = None
+    session_state["inference_state"] = None
+    return (
+        None,
+        gr.update(open=True),
+        None,
+        None,
+        gr.update(value=None, visible=False),
+        session_state,
+    )
+
+def clear_points(session_state):
+    session_state["input_points"] = []
+    session_state["input_labels"] = []
+    if session_state["inference_state"] is not None and session_state["inference_state"].get("tracking_has_started", False):
+        predictor.reset_state(session_state["inference_state"])
+    return (
+        session_state["first_frame"],
+        None,
+        gr.update(value=None, visible=False),
+        session_state,
+    )
 
-def clear_points(session):
-    session["input_points"] = []
-    session["input_labels"] = []
-    if session["inference_state"] and session["inference_state"].get("tracking_has_started"):
-        predictor.reset_state(session["inference_state"])
-    return session["first_frame"], None, gr.update(value=None, visible=False), session
+def preprocess_video_in(video_path, session_state):
+    if video_path is None:
+        return (
+            gr.update(open=True),
+            None,
+            None,
+            gr.update(value=None, visible=False),
+            session_state,
+        )
 
-def preprocess_video(video_path, session):
     cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened(): return gr.update(open=True), None, None, gr.update(value=None, visible=False), session
-    
+    if not cap.isOpened():
+        print("Error: Could not open video.")
+        return (
+            gr.update(open=True),
+            None,
+            None,
+            gr.update(value=None, visible=False),
+            session_state,
+        )
+
+    # Video properties
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    stride = max(1, total_frames // 300)
-    frames, first_frame = [], None
-    
-    w, h = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    target_w = 640
-    scale = target_w / w if w > target_w else 1.0
 
-    frame_id = 0
+    # Resize for CPU performance
+    target_width = 640
+    scale_factor = 1.0
+    if frame_width > target_width:
+        scale_factor = target_width / frame_width
+        frame_width = target_width
+        frame_height = int(frame_height * scale_factor)
+
+    # Read frames with stride for CPU optimization
+    frame_number = 0
+    first_frame = None
+    all_frames = []
+    frame_stride = max(1, total_frames // 300)  # Limit to ~300 frames
+
     while True:
         ret, frame = cap.read()
-        if not ret: break
-        if frame_id % stride == 0:
-            if scale < 1.0:
-                frame = cv2.resize(frame, (int(w*scale), int(h*scale)))
+        if not ret:
+            break
+        if frame_number % frame_stride == 0:
+            if scale_factor != 1.0:
+                frame = cv2.resize(frame, (frame_width, frame_height), interpolation=cv2.INTER_AREA)
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            if first_frame is None: first_frame = frame
-            frames.append(frame)
-        frame_id += 1
+            if first_frame is None:
+                first_frame = frame
+            all_frames.append(frame)
+        frame_number += 1
+
     cap.release()
+    session_state["first_frame"] = copy.deepcopy(first_frame)
+    session_state["all_frames"] = all_frames
+    session_state["frame_stride"] = frame_stride
+    session_state["scale_factor"] = scale_factor
+    session_state["original_dimensions"] = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
 
-    session.update({"first_frame": first_frame, "all_frames": frames, "frame_stride": stride, "scale_factor": scale, "inference_state": predictor.init_state(video_path=video_path), "input_points": [], "input_labels": []})
-    return gr.update(open=False), first_frame, None, gr.update(value=None, visible=False), session
-
-def show_mask(mask, obj_id=None):
-    cmap = plt.get_cmap("tab10")
-    color = np.array([*cmap(0 if obj_id is None else obj_id)[:3], 0.6])
-    h, w = mask.shape
-    mask_rgba = (mask.reshape(h, w, 1) * color.reshape(1, 1, -1) * 255).astype(np.uint8)
-    proper_mask = np.zeros((h, w, 4), dtype=np.uint8)
-    proper_mask[:, :, :min(mask_rgba.shape[2], 4)] = mask_rgba[:, :, :min(mask_rgba.shape[2], 4)]
-    return Image.fromarray(proper_mask, "RGBA")
-
-def segment_with_points(ptype, session, evt):
-    session["input_points"].append(evt.index)
-    session["input_labels"].append(1 if ptype == "include" else 0)
-    first = session["first_frame"]
-    h, w = first.shape[:2]
-    
-    layer = np.zeros((h, w, 4), dtype=np.uint8)
-    for idx, pt in enumerate(session["input_points"]):
-        color = (0, 255, 0, 255) if session["input_labels"][idx] == 1 else (255, 0, 0, 255)
-        cv2.circle(layer, pt, int(min(w, h)*0.01), color, -1)
-    
-    overlay = Image.alpha_composite(Image.fromarray(first).convert("RGBA"), Image.fromarray(layer, "RGBA"))
+    session_state["inference_state"] = predictor.init_state(video_path=video_path)
+    session_state["input_points"] = []
+    session_state["input_labels"] = []
+
+    return [
+        gr.update(open=False),
+        first_frame,
+        None,
+        gr.update(value=None, visible=False),
+        session_state,
+    ]
+
+def segment_with_points(point_type, session_state, evt: gr.SelectData):
+    session_state["input_points"].append(evt.index)
+    print(f"TRACKING INPUT POINT: {session_state['input_points']}")
+
+    if point_type == "include":
+        session_state["input_labels"].append(1)
+    elif point_type == "exclude":
+        session_state["input_labels"].append(0)
+    print(f"TRACKING INPUT LABEL: {session_state['input_labels']}")
+
+    first_frame = session_state["first_frame"]
+    h, w = first_frame.shape[:2]
+    transparent_background = Image.fromarray(first_frame).convert("RGBA")
+
+    # Draw points
+    fraction = 0.01
+    radius = int(fraction * min(w, h))
+    transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
+
+    for index, track in enumerate(session_state["input_points"]):
+        color = (0, 255, 0, 255) if session_state["input_labels"][index] == 1 else (255, 0, 0, 255)
+        cv2.circle(transparent_layer, track, radius, color, -1)
+
+    transparent_layer = Image.fromarray(transparent_layer, "RGBA")
+    selected_point_map = Image.alpha_composite(transparent_background, transparent_layer)
+
+    points = np.array(session_state["input_points"], dtype=np.float32)
+    labels = np.array(session_state["input_labels"], np.int32)
 
     try:
-        _, _, logits = predictor.add_new_points(session["inference_state"], 0, OBJ_ID, np.array(session["input_points"]), np.array(session["input_labels"]))
-        mask = (logits[0] > 0.0).cpu().numpy()
-        mask = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
-        mask_img = show_mask(mask)
-        return overlay, Image.alpha_composite(Image.fromarray(first).convert("RGBA"), mask_img), session
+        _, _, out_mask_logits = predictor.add_new_points(
+            inference_state=session_state["inference_state"],
+            frame_idx=0,
+            obj_id=OBJ_ID,
+            points=points,
+            labels=labels,
+        )
+        mask_array = (out_mask_logits[0] > 0.0).cpu().numpy()
+        
+        # Ensure mask matches frame size
+        if mask_array.shape[:2] != (h, w):
+            mask_array = cv2.resize(mask_array.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
+        
+        mask_image = show_mask(mask_array)
+        if mask_image.size != transparent_background.size:
+            mask_image = mask_image.resize(transparent_background.size, Image.NEAREST)
+        
+        first_frame_output = Image.alpha_composite(transparent_background, mask_image)
     except Exception as e:
-        print("Segmentation error:", e)
-        return overlay, overlay, session
+        print(f"Error in segmentation: {e}")
+        first_frame_output = selected_point_map
 
-def propagate(video_in, session, progress=gr.Progress()):
-    if not session["input_points"] or not session["inference_state"]: return None, session
+    return selected_point_map, first_frame_output, session_state
 
-    masks = {}
-    for i, (idxs, obj_ids, logits) in enumerate(predictor.propagate_in_video(session["inference_state"])):
-        try:
-            masks[idxs] = {oid: (logits[j] > 0.0).cpu().numpy() for j, oid in enumerate(obj_ids)}
-            progress(i / 300, desc=f"Tracking frame {idxs}")
-        except: continue
+def show_mask(mask, obj_id=None, random_color=False, convert_to_image=True):
+    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])
 
-    frames_out, stride = [], max(1, len(masks) // 50)
-    for i in range(0, len(masks), stride):
-        if i not in masks or OBJ_ID not in masks[i]: continue
+    h, w = mask.shape[-2:] if len(mask.shape) > 2 else mask.shape
+    mask_reshaped = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    mask_rgba = (mask_reshaped * 255).astype(np.uint8)
+
+    if convert_to_image:
         try:
-            frame = session["all_frames"][i]
-            mask = masks[i][OBJ_ID]
-            h, w = frame.shape[:2]
-            mask = cv2.resize(mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
-            output = Image.alpha_composite(Image.fromarray(frame).convert("RGBA"), show_mask(mask))
-            frames_out.append(np.array(output))
-        except: continue
-
-    out_path = os.path.join(tempfile.gettempdir(), f"output_video_{datetime.now().strftime('%Y%m%d%H%M%S')}.mp4")
-    fps = min(15, get_fps(video_in))
-    ImageSequenceClip(frames_out, fps=fps).write_videofile(out_path, codec="libx264", bitrate="800k", threads=2, logger=None)
-    gc.collect()
-    return gr.update(value=out_path, visible=True), session
+            if mask_rgba.shape[2] != 4:
+                proper_mask = np.zeros((h, w, 4), dtype=np.uint8)
+                proper_mask[:, :, :min(mask_rgba.shape[2], 4)] = mask_rgba[:, :, :min(mask_rgba.shape[2], 4)]
+                mask_rgba = proper_mask
+            return Image.fromarray(mask_rgba, "RGBA")
+        except Exception as e:
+            print(f"Error converting mask to image: {e}")
+            return Image.fromarray(np.zeros((h, w, 4), dtype=np.uint8), "RGBA")
+    
+    return mask_rgba
+
+def propagate_to_all(video_in, session_state, progress=gr.Progress()):
+    if len(session_state["input_points"]) == 0 or video_in is None or session_state["inference_state"] is None:
+        return gr.update(value=None, visible=False), session_state
+
+    chunk_size = 3
+    try:
+        video_segments = {}
+        total_frames = len(session_state["all_frames"])
+        progress(0, desc="Propagating segmentation through video...")
+
+        for i, (out_frame_idx, out_obj_ids, out_mask_logit) in enumerate(predictor.propagate_in_video(session_state["inference_state"])):
+            try:
+                video_segments[out_frame_idx] = {
+                    out_obj_id: (out_mask_logit[i] > 0.0).cpu().numpy()
+                    for i, out_obj_id in enumerate(out_obj_ids)
+                }
+                progress((i + 1) / total_frames, desc=f"Processed frame {out_frame_idx}/{total_frames}")
+                if out_frame_idx % chunk_size == 0:
+                    del out_mask_logit
+                    import gc
+                    gc.collect()
+            except Exception as e:
+                print(f"Error processing frame {out_frame_idx}: {e}")
+                continue
+
+        max_output_frames = 50
+        vis_frame_stride = max(1, total_frames // max_output_frames)
+        first_frame = session_state["all_frames"][0]
+        h, w = first_frame.shape[:2]
+        output_frames = []
+
+        for out_frame_idx in range(0, total_frames, vis_frame_stride):
+            if out_frame_idx not in video_segments or OBJ_ID not in video_segments[out_frame_idx]:
+                continue
+            try:
+                frame = session_state["all_frames"][out_frame_idx]
+                transparent_background = Image.fromarray(frame).convert("RGBA")
+                out_mask = video_segments[out_frame_idx][OBJ_ID]
+
+                # Validate mask dimensions
+                if out_mask.shape[:2] != (h, w):
+                    if out_mask.size == 0:  # Skip empty masks
+                        print(f"Skipping empty mask for frame {out_frame_idx}")
+                        continue
+                    out_mask = cv2.resize(out_mask.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
+
+                mask_image = show_mask(out_mask)
+                if mask_image.size != transparent_background.size:
+                    mask_image = mask_image.resize(transparent_background.size, Image.NEAREST)
+
+                output_frame = Image.alpha_composite(transparent_background, mask_image)
+                output_frames.append(np.array(output_frame))
+
+                if len(output_frames) % 10 == 0:
+                    import gc
+                    gc.collect()
+            except Exception as e:
+                print(f"Error creating output frame {out_frame_idx}: {e_RAW
+                traceback.print_exc()
+                continue
 
+        original_fps = get_video_fps(video_in)
+        fps = min(original_fps, 15)  # Cap at 15 FPS for CPU
+
+        clip = ImageSequenceClip(output_frames, fps=fps)
+        unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
+        final_vid_output_path = os.path.join(tempfile.gettempdir(), f"output_video_{unique_id}.mp4")
+
+        clip.write_videofile(
+            final_vid_output_path,
+            codec="libx264",
+            bitrate="800k",
+            threads=2,
+            logger=None
+        )
+
+        del video_segments, output_frames
+        import gc
+        gc.collect()
+
+        return gr.update(value=final_vid_output_path, visible=True), session_state
+
+    except Exception as e:
+        print(f"Error in propagate_to_all: {e}")
+        return gr.update(value=None, visible=False), session_state
+
+def update_ui():
+    return gr.update(visible=True)
+
+# Gradio Interface
 with gr.Blocks() as demo:
-    state = gr.State({"first_frame": None, "all_frames": None, "input_points": [], "input_labels": [], "inference_state": None, "frame_stride": 1, "scale_factor": 1.0, "original_dimensions": None})
-
-    gr.Markdown("<center><strong><font size='8'>EdgeTAM CPU</font></strong> <a href='https://github.com/facebookresearch/EdgeTAM'><font size='6'>[GitHub]</font></a></center>")
-
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("""<ol><li>Upload a video or use an example</li><li>Select 'include' or 'exclude' and click points</li><li>Click 'Track' to segment and track</li></ol>""")
-            drawer = gr.Accordion("Input Video", open=True)
-            with drawer:
-                video_in = gr.Video(label="Input Video", format="mp4")
-            ptype = gr.Radio(label="Point Type", choices=["include", "exclude"], value="include")
-            track_btn = gr.Button("Track", variant="primary")
-            clear_btn = gr.Button("Clear Points")
-            reset_btn = gr.Button("Reset")
-            points_map = gr.Image(label="Frame with Points", type="numpy", interactive=False)
-        with gr.Column():
-            gr.Markdown("# Try some examples ⬇️")
-            gr.Examples(examples, inputs=[video_in], examples_per_page=5)
-            output_img = gr.Image(label="Reference Mask")
-            output_vid = gr.Video(visible=False)
-
-    video_in.upload(preprocess_video, [video_in, state], [drawer, points_map, output_img, output_vid, state])
-    video_in.change(preprocess_video, [video_in, state], [drawer, points_map, output_img, output_vid, state])
-    points_map.select(segment_with_points, [ptype, state], [points_map, output_img, state])
-    clear_btn.click(clear_points, state, [points_map, output_img, output_vid, state])
-    reset_btn.click(reset, state, [video_in, drawer, points_map, output_img, output_vid, state])
-    track_btn.click(fn=propagate, inputs=[video_in, state], outputs=[output_vid, state])
-
-if __name__ == '__main__':
-    demo.queue()
-    demo.launch()
+    session_state = gr.State({
+        "first_frame": None,
+        "all_frames": None,
+        "input_points": [],
+        "input_labels": [],
+        "inference_state": None,
+        "frame_stride": 1,
+        "scale_factor": 1.0,
+        "original_dimensions": None,
+    })
+
+    with gr.Column():
+        gr.Markdown(title)
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown(description_p)
+                with gr.Accordion("Input Video", open=True) as video_in_drawer:
+                    video_in = gr.Video(label="Input Video", format="mp4")
+                with gr.Row():
+                    point_type = gr.Radio(label="point type", choices=["include", "exclude"], value="include", scale=2)
+                    propagate_btn = gr.Button("Track", scale=1, variant="primary")
+                    clear_points_btn = gr.Button("Clear Points", scale=1)
+                    reset_btn = gr.Button("Reset", scale=1)
+                points_map = gr.Image(label="Frame with Point Prompt", type="numpy", interactive=False)
+            with gr.Column():
+                gr.Markdown("# Try some of the examples below ⬇️")
+                gr.Examples(examples=examples, inputs=[video_in], examples_per_page=5)
+                output_image = gr.Image(label="Reference Mask")
+                output_video = gr.Video(visible=False)
+
+    video_in.upload(
+        fn=preprocess_video_in,
+        inputs=[video_in, session_state],
+        outputs=[video_in_drawer, points_map, output_image, output_video, session_state],
+        queue=False,
+    )
+
+    video_in.change(
+        fn=preprocess_video_in,
+        inputs=[video_in, session_state],
+        outputs=[video_in_drawer, points_map, output_image, output_video, session_state],
+        queue=False,
+    )
+
+    points_map.select(
+        fn=segment_with_points,
+        inputs=[point_type, session_state],
+        outputs=[points_map, output_image, session_state],
+        queue=False,
+    )
+
+    clear_points_btn.click(
+        fn=clear_points,
+        inputs=session_state,
+        outputs=[points_map, output_image, output_video, session_state],
+        queue=False,
+    )
+
+    reset_btn.click(
+        fn=reset,
+        inputs=session_state,
+        outputs=[video_in, video_in_drawer, points_map, output_image, output_video, session_state],
+        queue=False,
+    )
+
+    propagate_btn.click(
+        fn=update_ui,
+        inputs=[],
+        outputs=output_video,
+        queue=False,
+    ).then(
+        fn=propagate_to_all,
+        inputs=[video_in, session_state],
+        outputs=[output_video, session_state],
+        queue=True,
+    )
+
+demo.queue()
+demo.launch()