update sam2

app.py

@@ -7,15 +7,18 @@ import timm
 print("installed", timm.__version__)
 import gradio as gr
 from inference import sam_preprocess, beit3_preprocess
-from model.evf_sam import EvfSamModel
+from model.evf_sam2 import EvfSam2Model
+from model.evf_sam2_video import EvfSam2Model as EvfSam2VideoModel
 from transformers import AutoTokenizer
 import torch
+import cv2
 import numpy as np
 import sys
 import os
+import tqdm
 
-version = "YxZhang/evf-sam"
-model_type = "ori"
+version = "YxZhang/evf-sam2"
+model_type = "sam2"
 
 tokenizer = AutoTokenizer.from_pretrained(
     version,
@@ -26,27 +29,40 @@ tokenizer = AutoTokenizer.from_pretrained(
 kwargs = {
     "torch_dtype": torch.half,
 }
-model = EvfSamModel.from_pretrained(version, low_cpu_mem_usage=True,
-                                    **kwargs).eval()
-model.to('cuda')
+
+image_model = EvfSam2Model.from_pretrained(version,
+                                           low_cpu_mem_usage=True,
+                                           **kwargs)
+del image_model.visual_model.memory_encoder
+del image_model.visual_model.memory_attention
+image_model = image_model.eval()
+image_model.to('cuda')
+
+video_model = EvfSam2VideoModel.from_pretrained(version,
+                                                low_cpu_mem_usage=True,
+                                                **kwargs)
+video_model = video_model.eval()
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+video_model.to('cuda')
 
 
 @spaces.GPU
 @torch.no_grad()
-def pred(image_np, prompt):
+def inference_image(image_np, prompt):
     original_size_list = [image_np.shape[:2]]
 
-    image_beit = beit3_preprocess(image_np, 224).to(dtype=model.dtype,
-                                                    device=model.device)
+    image_beit = beit3_preprocess(image_np, 224).to(dtype=image_model.dtype,
+                                                    device=image_model.device)
 
     image_sam, resize_shape = sam_preprocess(image_np, model_type=model_type)
-    image_sam = image_sam.to(dtype=model.dtype, device=model.device)
+    image_sam = image_sam.to(dtype=image_model.dtype,
+                             device=image_model.device)
 
     input_ids = tokenizer(
-        prompt, return_tensors="pt")["input_ids"].to(device=model.device)
+        prompt, return_tensors="pt")["input_ids"].to(device=image_model.device)
 
     # infer
-    pred_mask = model.inference(
+    pred_mask = image_model.inference(
         image_sam.unsqueeze(0),
         image_beit.unsqueeze(0),
         input_ids,
@@ -61,7 +77,50 @@ def pred(image_np, prompt):
                                 pred_mask[:, :, None].astype(np.uint8) *
                                 np.array([50, 120, 220]) * 0.5)[pred_mask]
 
-    return visualization / 255.0, pred_mask.astype(np.float16)
+    return visualization / 255.0
+
+
+@spaces.GPU
+@torch.no_grad()
+@torch.autocast(device_type="cuda", dtype=torch.float16)
+def inference_video(video_path, prompt):
+
+    os.system("rm -rf demo_temp")
+    os.makedirs("demo_temp/input_frames", exist_ok=True)
+    os.system(
+        "ffmpeg -i {} -q:v 2 -start_number 0 demo_temp/input_frames/'%05d.jpg'"
+        .format(video_path))
+    input_frames = sorted(os.listdir("demo_temp/input_frames"))
+    image_np = cv2.imread("demo_temp/input_frames/00000.jpg")
+    image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
+
+    height, width, channels = image_np.shape
+
+    image_beit = beit3_preprocess(image_np, 224).to(dtype=video_model.dtype,
+                                                    device=video_model.device)
+
+    input_ids = tokenizer(
+        prompt, return_tensors="pt")["input_ids"].to(device=video_model.device)
+
+    # infer
+    output = video_model.inference(
+        "demo_temp/input_frames",
+        image_beit.unsqueeze(0),
+        input_ids,
+    )
+    # save visualization
+    video_writer = cv2.VideoWriter("demo_temp/out.mp4", fourcc, 30,
+                                   (width, height))
+    pbar = tqdm(input_frames)
+    pbar.set_description("generating video: ")
+    for i, file in enumerate(pbar):
+        img = cv2.imread(os.path.join("demo_temp/input_frames", file))
+        vis = img + np.array([0, 0, 128]) * output[i][1].transpose(1, 2, 0)
+        vis = np.clip(vis, 0, 255)
+        vis = np.uint8(vis)
+        video_writer.write(vis)
+    video_writer.release()
+    return "demo_temp/out.mp4"
 
 
 desc = """
@@ -73,28 +132,41 @@ desc = """
 # desc_title_str = '<div align ="center"><img src="assets/logo.jpg" width="20%"><h3> Early Vision-Language Fusion for Text-Prompted Segment Anything Model</h3></div>'
 # desc_link_str = '[![arxiv paper](https://img.shields.io/badge/arXiv-Paper-red)](https://arxiv.org/abs/2406.20076)'
 
-demo = gr.Interface(
-    fn=pred,
-    inputs=[
-        gr.components.Image(type="numpy", label="Image", image_mode="RGB"),
-        gr.components.Textbox(
-            label="Prompt",
-            info=
-            "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
-        )
-    ],
-    outputs=[
-        gr.components.Image(type="numpy", label="visulization"),
-        gr.components.Image(type="numpy", label="mask")
-    ],
-    examples=[["assets/zebra.jpg", "zebra top left"],
-              ["assets/bus.jpg", "bus going to south common"],
-              [
-                  "assets/carrots.jpg",
-                  "3carrots in center with ice and greenn leaves"
-              ]],
-    title="📷 EVF-SAM: Referring Expression Segmentation",
-    description=desc,
-    allow_flagging="never")
-# demo.launch()
-demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown(desc)
+    with gr.Tab(label="EVF-SAM-2-Image"):
+        with gr.Row():
+            input_image = gr.Image(type='numpy',
+                                   label='Input Image',
+                                   image_mode='RGB')
+            output_image = gr.Image(type='numpy', label='Output Image')
+        with gr.Row():
+            image_prompt = gr.Textbox(
+                label="Prompt",
+                info=
+                "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
+            )
+            submit_image = gr.Button(value='Submit',
+                                     scale=1,
+                                     variant='primary')
+    with gr.Tab(label="EVF-SAM-2-Video"):
+        with gr.Row():
+            input_video = gr.Video(label='Input Video')
+            output_video = gr.Video(label='Output Video')
+        with gr.Row():
+            video_prompt = gr.Textbox(
+                label="Prompt",
+                info=
+                "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
+            )
+            submit_video = gr.Button(value='Submit',
+                                     scale=1,
+                                     variant='primary')
+
+    submit_image.click(fn=inference_image,
+                       inputs=[input_image, image_prompt],
+                       outputs=output_image)
+    submit_video.click(fn=inference_video,
+                       inputs=[input_video, video_prompt],
+                       outputs=output_video)
+demo.launch(show_error=True)

app_video.py → app.py.bak

@@ -7,18 +7,15 @@ import timm
 print("installed", timm.__version__)
 import gradio as gr
 from inference import sam_preprocess, beit3_preprocess
-from model.evf_sam2 import EvfSam2Model
-from model.evf_sam2_video import EvfSam2Model as EvfSam2VideoModel
+from model.evf_sam import EvfSamModel
 from transformers import AutoTokenizer
 import torch
-import cv2
 import numpy as np
 import sys
 import os
-import tqdm
 
-version = "YxZhang/evf-sam2"
-model_type = "sam2"
+version = "YxZhang/evf-sam"
+model_type = "ori"
 
 tokenizer = AutoTokenizer.from_pretrained(
     version,
@@ -29,40 +26,27 @@ tokenizer = AutoTokenizer.from_pretrained(
 kwargs = {
     "torch_dtype": torch.half,
 }
-
-image_model = EvfSam2Model.from_pretrained(version,
-                                           low_cpu_mem_usage=True,
-                                           **kwargs)
-del image_model.visual_model.memory_encoder
-del image_model.visual_model.memory_attention
-image_model = image_model.eval()
-image_model.to('cuda')
-
-video_model = EvfSam2VideoModel.from_pretrained(version,
-                                                low_cpu_mem_usage=True,
-                                                **kwargs)
-video_model = video_model.eval()
-fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-video_model.to('cuda')
+model = EvfSamModel.from_pretrained(version, low_cpu_mem_usage=True,
+                                    **kwargs).eval()
+model.to('cuda')
 
 
 @spaces.GPU
 @torch.no_grad()
-def inference_image(image_np, prompt):
+def pred(image_np, prompt):
     original_size_list = [image_np.shape[:2]]
 
-    image_beit = beit3_preprocess(image_np, 224).to(dtype=image_model.dtype,
-                                                    device=image_model.device)
+    image_beit = beit3_preprocess(image_np, 224).to(dtype=model.dtype,
+                                                    device=model.device)
 
     image_sam, resize_shape = sam_preprocess(image_np, model_type=model_type)
-    image_sam = image_sam.to(dtype=image_model.dtype,
-                             device=image_model.device)
+    image_sam = image_sam.to(dtype=model.dtype, device=model.device)
 
     input_ids = tokenizer(
-        prompt, return_tensors="pt")["input_ids"].to(device=image_model.device)
+        prompt, return_tensors="pt")["input_ids"].to(device=model.device)
 
     # infer
-    pred_mask = image_model.inference(
+    pred_mask = model.inference(
         image_sam.unsqueeze(0),
         image_beit.unsqueeze(0),
         input_ids,
@@ -77,50 +61,7 @@ def inference_image(image_np, prompt):
                                 pred_mask[:, :, None].astype(np.uint8) *
                                 np.array([50, 120, 220]) * 0.5)[pred_mask]
 
-    return visualization / 255.0
-
-
-@spaces.GPU
-@torch.no_grad()
-@torch.autocast(device_type="cuda", dtype=torch.float16)
-def inference_video(video_path, prompt):
-
-    os.system("rm -rf demo_temp")
-    os.makedirs("demo_temp/input_frames", exist_ok=True)
-    os.system(
-        "ffmpeg -i {} -q:v 2 -start_number 0 demo_temp/input_frames/'%05d.jpg'"
-        .format(video_path))
-    input_frames = sorted(os.listdir("demo_temp/input_frames"))
-    image_np = cv2.imread("demo_temp/input_frames/00000.jpg")
-    image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
-
-    height, width, channels = image_np.shape
-
-    image_beit = beit3_preprocess(image_np, 224).to(dtype=video_model.dtype,
-                                                    device=video_model.device)
-
-    input_ids = tokenizer(
-        prompt, return_tensors="pt")["input_ids"].to(device=video_model.device)
-
-    # infer
-    output = video_model.inference(
-        "demo_temp/input_frames",
-        image_beit.unsqueeze(0),
-        input_ids,
-    )
-    # save visualization
-    video_writer = cv2.VideoWriter("demo_temp/out.mp4", fourcc, 30,
-                                   (width, height))
-    pbar = tqdm(input_frames)
-    pbar.set_description("generating video: ")
-    for i, file in enumerate(pbar):
-        img = cv2.imread(os.path.join("demo_temp/input_frames", file))
-        vis = img + np.array([0, 0, 128]) * output[i][1].transpose(1, 2, 0)
-        vis = np.clip(vis, 0, 255)
-        vis = np.uint8(vis)
-        video_writer.write(vis)
-    video_writer.release()
-    return "demo_temp/out.mp4"
+    return visualization / 255.0, pred_mask.astype(np.float16)
 
 
 desc = """
@@ -132,41 +73,28 @@ desc = """
 # desc_title_str = '<div align ="center"><img src="assets/logo.jpg" width="20%"><h3> Early Vision-Language Fusion for Text-Prompted Segment Anything Model</h3></div>'
 # desc_link_str = '[![arxiv paper](https://img.shields.io/badge/arXiv-Paper-red)](https://arxiv.org/abs/2406.20076)'
 
-with gr.Blocks() as demo:
-    gr.Markdown(desc)
-    with gr.Tab(label="EVF-SAM-2-Image"):
-        with gr.Row():
-            input_image = gr.Image(type='numpy',
-                                   label='Input Image',
-                                   image_mode='RGB')
-            output_image = gr.Image(type='numpy', label='Output Image')
-        with gr.Row():
-            image_prompt = gr.Textbox(
-                label="Prompt",
-                info=
-                "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
-            )
-            submit_image = gr.Button(value='Submit',
-                                     scale=1,
-                                     variant='primary')
-    with gr.Tab(label="EVF-SAM-2-Video"):
-        with gr.Row():
-            input_video = gr.Video(label='Input Video')
-            output_video = gr.Video(label='Output Video')
-        with gr.Row():
-            video_prompt = gr.Textbox(
-                label="Prompt",
-                info=
-                "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
-            )
-            submit_video = gr.Button(value='Submit',
-                                     scale=1,
-                                     variant='primary')
-
-    submit_image.click(fn=inference_image,
-                       inputs=[input_image, image_prompt],
-                       outputs=output_image)
-    submit_video.click(fn=inference_video,
-                       inputs=[input_video, video_prompt],
-                       outputs=output_video)
-demo.launch(show_error=True)
+demo = gr.Interface(
+    fn=pred,
+    inputs=[
+        gr.components.Image(type="numpy", label="Image", image_mode="RGB"),
+        gr.components.Textbox(
+            label="Prompt",
+            info=
+            "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
+        )
+    ],
+    outputs=[
+        gr.components.Image(type="numpy", label="visulization"),
+        gr.components.Image(type="numpy", label="mask")
+    ],
+    examples=[["assets/zebra.jpg", "zebra top left"],
+              ["assets/bus.jpg", "bus going to south common"],
+              [
+                  "assets/carrots.jpg",
+                  "3carrots in center with ice and greenn leaves"
+              ]],
+    title="📷 EVF-SAM: Referring Expression Segmentation",
+    description=desc,
+    allow_flagging="never")
+# demo.launch()
+demo.launch()