File size: 4,181 Bytes
032e687 |
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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
# Copyright (c) OpenMMLab. All rights reserved.
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_8b/'
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()
|