Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,212 Bytes
d59f323 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import argparse
import os
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer
import cv2
try:
from mmengine.visualization import Visualizer
except ImportError:
Visualizer = None
print("Warning: mmengine is not installed, visualization is disabled.")
def parse_args():
parser = argparse.ArgumentParser(description='Video Reasoning Segmentation')
parser.add_argument('image_folder', help='Path to image file')
parser.add_argument('--model_path', default="ByteDance/Sa2VA-8B")
parser.add_argument('--work-dir', default=None, help='The dir to save results.')
parser.add_argument('--text', type=str, default="<image>Please describe the video content.")
parser.add_argument('--select', type=int, default=-1)
args = parser.parse_args()
return args
def visualize(pred_mask, image_path, work_dir):
visualizer = Visualizer()
img = cv2.imread(image_path)
visualizer.set_image(img)
visualizer.draw_binary_masks(pred_mask, colors='g', alphas=0.4)
visual_result = visualizer.get_image()
output_path = os.path.join(work_dir, os.path.basename(image_path))
cv2.imwrite(output_path, visual_result)
if __name__ == "__main__":
cfg = parse_args()
model_path = cfg.model_path
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype="auto",
device_map="auto",
trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
model_path,
trust_remote_code=True
)
image_files = []
image_paths = []
image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".tiff"}
for filename in sorted(list(os.listdir(cfg.image_folder))):
if os.path.splitext(filename)[1].lower() in image_extensions:
image_files.append(filename)
image_paths.append(os.path.join(cfg.image_folder, filename))
vid_frames = []
for img_path in image_paths:
img = Image.open(img_path).convert('RGB')
vid_frames.append(img)
if cfg.select > 0:
img_frame = vid_frames[cfg.select - 1]
print(f"Selected frame {cfg.select}")
print(f"The input is:\n{cfg.text}")
result = model.predict_forward(
image=img_frame,
text=cfg.text,
tokenizer=tokenizer,
)
else:
print(f"The input is:\n{cfg.text}")
result = model.predict_forward(
video=vid_frames,
text=cfg.text,
tokenizer=tokenizer,
)
prediction = result['prediction']
print(f"The output is:\n{prediction}")
if '[SEG]' in prediction and Visualizer is not None:
_seg_idx = 0
pred_masks = result['prediction_masks'][_seg_idx]
for frame_idx in range(len(vid_frames)):
pred_mask = pred_masks[frame_idx]
if cfg.work_dir:
os.makedirs(cfg.work_dir, exist_ok=True)
visualize(pred_mask, image_paths[frame_idx], cfg.work_dir)
else:
os.makedirs('./temp_visualize_results', exist_ok=True)
visualize(pred_mask, image_paths[frame_idx], './temp_visualize_results')
else:
pass
|