File size: 4,119 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 123 124 125 126 127 128 129 130 131 |
# Copyright (c) OpenMMLab. All rights reserved.
import math
import os
import json
import torch
from transformers import AutoTokenizer, AutoModel
import cv2
from pycocotools import mask as maskUtils
from PIL import Image
images_folder = '/mnt/bn/xiangtai-training-data/project/xiangtai-windows/tt_vlm/qilu_folder/images/'
save_demo_dir = '/mnt/bn/xiangtai-training-data/project/xiangtai-windows/tt_vlm/qilu_folder/draw_images/'
save_mask_dir = '/mnt/bn/xiangtai-training-data/project/xiangtai-windows/tt_vlm/qilu_folder/masks/'
demo_items = []
image_files = os.listdir(images_folder)
gcg_question = '<image>\nCan you provide a brief description of the this image? Please output with interleaved segmentation masks for the corresponding phrases.'
for image_file in image_files:
image_path = os.path.join(images_folder, image_file)
demo_items.append(
{'image_path': image_path, 'question': gcg_question}
)
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,
# )
json_dir = '/mnt/bn/xiangtai-training-data/project/xiangtai-windows/glamm/groundingLMM/'
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')
json_path = os.path.join(json_dir, f'0{i+1}.json')
with open(json_path, 'r') as f:
results = json.load(f)
masks = results['pred_masks']
pred_masks = [torch.Tensor(maskUtils.decode(ann)) for ann in masks]
show_mask_pred(ori_image, pred_masks,
save_dir_demo=os.path.join(save_demo_dir, f'output_{i}.png'),
save_dir_mask=os.path.join(save_mask_dir, f"output_{i}.png")
)
def show_mask_pred(image, masks, save_dir_demo=None, save_dir_mask=None):
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()
_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]
mask_image = np.array(_mask_image).astype(np.uint8)
mask_image = Image.fromarray(mask_image)
mask_image.save(save_dir_mask)
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_demo)
return
if __name__ == '__main__':
main()
|