Upload app.py

app.py

@@ -1,39 +1,35 @@
 import gradio as gr
+import os
+import warnings
+
+so_path = "models/GroundingDINO/ops/MultiScaleDeformableAttention.cpython-39-x86_64-linux-gnu.so"
+if not os.path.exists(so_path):
+    os.system("python models/GroundingDINO/ops/setup.py build_ext develop --user")
+
 import torchvision.transforms as T
 from models import build_model
-import os
 import torch
 import misc as utils
 import numpy as np
 import torch.nn.functional as F
-import matplotlib.colors
 from torchvision.io import read_video
 import torchvision.transforms.functional as Func
 from ruamel.yaml import YAML
 from easydict import EasyDict
 from misc import nested_tensor_from_videos_list
 from torch.cuda.amp import autocast
-from PIL import Image, ImageDraw, ImageFont
-from rich.progress import track
+from PIL import Image, ImageDraw
 import imageio.v3 as iio
-import imageio
 import cv2
-import warnings
 import tempfile
 import argparse
 import time
 from huggingface_hub import hf_hub_download
 
-warnings.filterwarnings("ignore")
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# TARGET_FPS = 10
-
-so_path = "models/GroundingDINO/ops/MultiScaleDeformableAttention.cpython-39-x86_64-linux-gnu.so"
-if not os.path.exists(so_path):
-    # print("Building MultiScaleDeformableAttention module...")
-    # subprocess.run(["python", "setup.py", "build_ext", "--inplace"], cwd="models/GroundingDINO/ops", check=True)
-    os.system("python models/GroundingDINO/ops/setup.py build_ext develop --user")
+DURATION = 6
+CHECKPOINT = "ryb_mevis_swinb.pth"
 
 # Transform for video frames
 transform = T.Compose([
@@ -50,7 +46,7 @@ color_list = color_list.astype('uint8').tolist()
 model = None
 
 
-def load_model_once(config_path, checkpoint_path, device='cpu'):
+def load_model_once(config_path, device='cpu'):
     """Load model once at startup"""
     global model
     if model is None:
@@ -62,19 +58,12 @@ def load_model_once(config_path, checkpoint_path, device='cpu'):
 
         args = EasyDict(config)
         args.device = device
-        args.checkpoint_path = checkpoint_path
 
         model = build_model(args)
         model.to(device)
-
-        ckpt_path = hf_hub_download(
-            repo_id="Birdy80/888",
-            filename="ckpt/ryb_mevis_swinb.pth",
-            token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
-            repo_type="model"
-        )
-        
-        checkpoint = torch.load(ckpt_path, map_location='cpu')
+        cache_file = hf_hub_download(repo_id="liangtm/referdino", filename=CHECKPOINT)
+        # cache_file = 'ckpt/' + CHECKPOINT
+        checkpoint = torch.load(cache_file, map_location='cpu')
         state_dict = checkpoint["model_state_dict"]
         model.load_state_dict(state_dict, strict=False)
         model.eval()
@@ -118,13 +107,12 @@ def vis_add_mask(img, mask, color, edge_width=3):
     return origin_img
 
 
-def run_video_inference(input_video, text_prompt, tracking_alpha=0.1, target_fps=10):
+def run_video_inference(input_video, text_prompt, tracking_alpha=0.1, fps=15):
     """Main inference function for Gradio"""
     global model
     model.tracking_alpha = tracking_alpha
 
     # Set default values for other parameters
-    frame_step = 1
     show_box = True
     mask_edge_width = 6
 
@@ -137,17 +125,16 @@ def run_video_inference(input_video, text_prompt, tracking_alpha=0.1, target_fps
     # Process text prompt
     exp = " ".join(text_prompt.lower().split())
 
-    # Read video (≤8s on CPU)
-    video_frames, _, info = read_video(input_video, end_pts=8, pts_unit='sec')  # (T, H, W, C)
-    original_fps = info['video_fps']
-    frame_interval = max(round(original_fps / target_fps), 1)
+    # Read video
+    video_frames, _, info = read_video(input_video, end_pts=DURATION, pts_unit='sec')  # (T, H, W, C)
+
+    frame_step = max(round(info['video_fps'] / fps), 1)
 
     frames = []
-    for i in range(0, len(video_frames), frame_interval):
+    for i in range(0, len(video_frames), frame_step):
         source_frame = Func.to_pil_image(video_frames[i].permute(2, 0, 1))
         frames.append(source_frame)
 
-    info['video_fps'] = target_fps
     video_len = len(frames)
     if video_len == 0:
         return None, "No frames found in the video."
@@ -161,92 +148,60 @@ def run_video_inference(input_video, text_prompt, tracking_alpha=0.1, target_fps
 
     device = next(model.parameters()).device
     imgs = torch.stack(imgs, dim=0).to(device)
-    samples = nested_tensor_from_videos_list(imgs[None], size_divisibility=1)
+    samples = nested_tensor_from_videos_list(imgs[None], size_divisibility=16)
     img_h, img_w = imgs.shape[-2:]
     size = torch.as_tensor([int(img_h), int(img_w)]).to(device)
     target = {"size": size}
 
     start_infer = time.time()
+    # Run inference
     with torch.no_grad():
         with autocast(True):
             outputs = model(samples, [exp], [target])
     end_infer = time.time()
 
     pred_logits = outputs["pred_logits"][0]  # [t, q, k]
-    pred_masks = outputs["pred_masks"][0]    # [t, q, h, w]
-    pred_boxes = outputs["pred_boxes"][0]    # [t, q, 4]
+    pred_masks = outputs["pred_masks"][0]  # [t, q, h, w]
+    pred_boxes = outputs["pred_boxes"][0]  # [t, q, 4]
 
     # Select the query index according to pred_logits
-    pred_scores = pred_logits.sigmoid()      # [t, q, k]
-    pred_scores = pred_scores.mean(0)        # [q, K]
-    max_scores, _ = pred_scores.max(-1)      # [q,]
-    _, max_ind = max_scores.max(-1)          # [1,]
+    pred_scores = pred_logits.sigmoid()  # [t, q, k]
+    pred_scores = pred_scores.mean(0)  # [q, K]
+    max_scores, _ = pred_scores.max(-1)  # [q,]
+    _, max_ind = max_scores.max(-1)  # [1,]
     max_inds = max_ind.repeat(video_len)
     pred_masks = pred_masks[range(video_len), max_inds, ...]  # [t, h, w]
     pred_masks = pred_masks.unsqueeze(0)
     pred_boxes = pred_boxes[range(video_len), max_inds].cpu().numpy()  # [t, 4]
 
-    # Unpad and resize masks to original frame size
+    # Unpad and resize
     pred_masks = pred_masks[:, :, :img_h, :img_w].cpu()
     pred_masks = F.interpolate(pred_masks, size=(origin_h, origin_w), mode='bilinear', align_corners=False)
     pred_masks = (pred_masks.sigmoid() > 0.5).squeeze(0).cpu().numpy()
 
-    # Visualization color
-    color = "#DC143C"
-    color = (np.array(matplotlib.colors.hex2color(color)) * 255).astype('uint8')
+    # Visualization
+    color = np.array([220,  20,  60], dtype=np.uint8)
 
-    # ---- Save result video (with per-frame sanitization) ----
     start_save = time.time()
     save_imgs = []
-    for t, pil_img in enumerate(frames):
-        # draw mask
-        pil_img = vis_add_mask(pil_img, pred_masks[t], color, mask_edge_width)
+    for t, img in enumerate(frames):
+        # Draw mask
+        img = vis_add_mask(img, pred_masks[t], color, mask_edge_width)
 
-        # draw bbox
-        draw = ImageDraw.Draw(pil_img)
+        draw = ImageDraw.Draw(img)
         draw_boxes = pred_boxes[t][None]
         draw_boxes = rescale_bboxes(draw_boxes, (origin_w, origin_h)).tolist()
+
+        # Draw box if enabled
         if show_box:
             xmin, ymin, xmax, ymax = draw_boxes[0]
             draw.rectangle(((xmin, ymin), (xmax, ymax)), outline=tuple(color), width=5)
 
-        # --- Force sanitize to (H, W, 3) uint8 contiguous & even size ---
-        arr = np.asarray(pil_img)
-        if arr.ndim == 2:  # gray -> 3ch
-            arr = np.stack([arr] * 3, axis=-1)
-        elif arr.ndim == 3 and arr.shape[2] == 4:  # RGBA -> RGB
-            arr = arr[..., :3]
-        elif arr.ndim != 3:
-            arr = np.asarray(Image.fromarray(arr).convert("RGB"))
-
-        if arr.dtype != np.uint8:
-            arr = arr.astype(np.uint8)
+        save_imgs.append(np.asarray(img).copy())
 
-        h, w = arr.shape[:2]
-        if (h % 2) or (w % 2):  # x264 prefers even dims
-            arr = arr[: h - (h % 2), : w - (w % 2), :]
-
-        arr = np.ascontiguousarray(arr)
-        assert arr.ndim == 3 and arr.shape[2] == 3, f"Bad frame at t={t}, shape={arr.shape}, dtype={arr.dtype}"
-
-        save_imgs.append(arr)
-
-    # Write once (more robust than append_data)
+    # Save result video
     with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
-        fps = int(info['video_fps'] / frame_step)
-    
-        writer = imageio.get_writer(
-            tmp_file.name,
-            format='FFMPEG',
-            mode='I',
-            fps=fps,
-            codec='libx264',
-            output_params=['-pix_fmt', 'yuv420p'],
-            macro_block_size=None,
-        )
-        for arr in save_imgs:
-            writer.append_data(arr)
-        writer.close()
+        iio.imwrite(tmp_file.name, save_imgs, fps=fps)
         result_video_path = tmp_file.name
 
     end_save = time.time()
@@ -257,24 +212,19 @@ def run_video_inference(input_video, text_prompt, tracking_alpha=0.1, target_fps
     )
     return result_video_path, status
 
-    # except Exception as e:
-    #     return None, f"❌ Error during inference: {str(e)}"
-
 
 def main():
     # Configuration
     config_path = "configs/ytvos_swinb.yaml"  # Update this path
-    checkpoint_path = "ckpt/ryb_mevis_swinb.pth"  # Update this path
-    # device = "cuda" if torch.cuda.is_available() else "cpu"
-    device = "cpu"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "cpu"
 
     # Load model at startup
     print("Loading model...")
-    load_model_once(config_path, checkpoint_path, device)
+    load_model_once(config_path, device)
     print(f"Model loaded on device: {device}")
 
     # Create Gradio interface
-    # with gr.Blocks(title="ReferDINO") as demo:
     with gr.Blocks(
         title="ReferDINO",
         css="""
@@ -291,7 +241,7 @@ def main():
             <h3>Referring Video Object Segmentation with 
                 <a href="https://github.com/iSEE-Laboratory/ReferDINO">ReferDINO</a>
             </h3>
-            <h3>This is a CPU-only demo, so the video will be trimmed to ≤8 seconds.</h3>
+            <h3>Note that this demo runs on CPU, so the video will be trimmed to ≤6 seconds.</h3>
             """,
             elem_id="hero",
         )
@@ -323,15 +273,16 @@ def main():
                     maximum=1.0,
                     value=0.1,
                     step=0.05,
-                    info="Controls the memory updating (lower = longer memory)"
+                    info="controls the memory updating (lower = longer memory)"
                 )
+
                 target_fps = gr.Slider(
-                    label="Target FPS",
+                    label="FPS",
                     minimum=1,
-                    maximum=20,
+                    maximum=30,
                     value=10,
                     step=1,
-                    info="Downsample input video to this FPS for inference & output (Higher FPS will slow down processing)"
+                    info="controls the FPS (lower = faster processing)"
                 )
 
             with gr.Column(scale=1):
@@ -349,14 +300,14 @@ def main():
         # Examples
         gr.Examples(
             examples=[
-                ["example_video.mp4", "person walking", 0.1, 10],
-                ["example_video2.mp4", "red car", 0.1, 10],
+                ["dogs.mp4", "the dog drinking Sprite", 0.1, 10],
+                ["dogs.mp4", "the dog sleeping", 0.1, 10],
             ],
             inputs=[input_video, text_prompt, tracking_alpha, target_fps],
             outputs=[output_video],
             fn=run_video_inference,
             cache_examples=False,
-            label="Try these examples:"
+            label="📋 Try these examples:"
         )
 
         # Event handlers