|
|
|
|
|
import math |
|
|
import torch |
|
|
from transformers import AutoTokenizer, AutoModel |
|
|
import cv2 |
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
sam_prefix = '/mnt/bn/xiangtai-training-data-video/dataset/segmentation_datasets/sam_v_full/sav_000/sav_train/sav_000/' |
|
|
coco_prefix = './data/glamm_data/images/coco2014/train2014/' |
|
|
sam_p2 = 'data/sa_eval/' |
|
|
|
|
|
demo_items = [ |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nPlease describe the image.'}, |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nPlease segment the person.'}, |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nPlease segment the snowboard.'}, |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nPlease segment the snow.'}, |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nPlease segment the trees.'}, |
|
|
{'image_path': coco_prefix+'COCO_train2014_000000581921.jpg', 'question': '<image>\nCould you please give me a brief description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.'}, |
|
|
] |
|
|
|
|
|
TORCH_DTYPE_MAP = dict( |
|
|
fp16=torch.float16, bf16=torch.bfloat16, fp32=torch.float32, auto='auto') |
|
|
|
|
|
def get_video_frames(video_path): |
|
|
cap = cv2.VideoCapture(video_path) |
|
|
|
|
|
if not cap.isOpened(): |
|
|
print("Error: Cannot open video file.") |
|
|
return |
|
|
|
|
|
frames = [] |
|
|
|
|
|
frame_id = 0 |
|
|
while True: |
|
|
ret, frame = cap.read() |
|
|
|
|
|
if not ret: |
|
|
break |
|
|
|
|
|
frames.append(frame) |
|
|
|
|
|
frame_id += 1 |
|
|
|
|
|
cap.release() |
|
|
return frames |
|
|
|
|
|
def get_frames_from_video(video_path, n_frames=5): |
|
|
frames = get_video_frames(video_path) |
|
|
stride = len(frames) / (n_frames + 1e-4) |
|
|
ret = [] |
|
|
for i in range(n_frames): |
|
|
idx = int(i * stride) |
|
|
frame = frames[idx] |
|
|
frame = frame[:, :, ::-1] |
|
|
frame_image = Image.fromarray(frame).convert('RGB') |
|
|
ret.append(frame_image) |
|
|
return ret |
|
|
|
|
|
|
|
|
def main(): |
|
|
|
|
|
model_path = 'work_dirs/sa2va_4b/' |
|
|
model = AutoModel.from_pretrained( |
|
|
model_path, |
|
|
torch_dtype=torch.bfloat16, |
|
|
low_cpu_mem_usage=True, |
|
|
use_flash_attn=True, |
|
|
trust_remote_code=True, |
|
|
).eval().cuda() |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained( |
|
|
model_path, |
|
|
trust_remote_code=True, |
|
|
) |
|
|
|
|
|
for i, demo_item in enumerate(demo_items): |
|
|
image_path = demo_item['image_path'] |
|
|
text_prompts = demo_item['question'] |
|
|
ori_image = Image.open(image_path).convert('RGB') |
|
|
input_dict = { |
|
|
'image': ori_image, |
|
|
'text': text_prompts, |
|
|
'past_text': '', |
|
|
'mask_prompts': None, |
|
|
'tokenizer': tokenizer, |
|
|
} |
|
|
|
|
|
return_dict = model.predict_forward(**input_dict) |
|
|
print(i, ': ', return_dict['prediction']) |
|
|
if 'prediction_masks' in return_dict.keys() and return_dict['prediction_masks'] and len(return_dict['prediction_masks']) != 0: |
|
|
show_mask_pred(ori_image, return_dict['prediction_masks'], save_dir=f'./demos/output_{i}.png') |
|
|
|
|
|
def show_mask_pred(image, masks, save_dir='./output.png'): |
|
|
from PIL import Image |
|
|
import numpy as np |
|
|
|
|
|
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), |
|
|
(255, 255, 0), (255, 0, 255), (0, 255, 255), |
|
|
(128, 128, 255)] |
|
|
|
|
|
masks = torch.stack(masks, dim=0).cpu().numpy()[:, 0] |
|
|
_mask_image = np.zeros((masks.shape[1], masks.shape[2], 3), dtype=np.uint8) |
|
|
|
|
|
for i, mask in enumerate(masks): |
|
|
color = colors[i % len(colors)] |
|
|
_mask_image[:, :, 0] = _mask_image[:, :, 0] + mask.astype(np.uint8) * color[0] |
|
|
_mask_image[:, :, 1] = _mask_image[:, :, 1] + mask.astype(np.uint8) * color[1] |
|
|
_mask_image[:, :, 2] = _mask_image[:, :, 2] + mask.astype(np.uint8) * color[2] |
|
|
|
|
|
|
|
|
image = np.array(image) |
|
|
image = image * 0.5 + _mask_image * 0.5 |
|
|
image = image.astype(np.uint8) |
|
|
image = Image.fromarray(image) |
|
|
image.save(save_dir) |
|
|
|
|
|
return |
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
