File size: 7,400 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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 |
import os
from mmengine.dist import master_only
from vlm.datasets.evaluation.base_eval_dataset import BaseEvalDataset
import json
import numpy as np
import copy
import cv2
from PIL import Image
from lmdeploy.vl.constants import IMAGE_TOKEN
import pycocotools.mask as maskUtils
class SAM2Dataset(BaseEvalDataset):
METAINFO: dict = dict(name='image dataset')
def __init__(
self,
video_folder,
json_folder,
bs=8,
select_frames=3,
):
super().__init__()
self.json_folder = json_folder
json_files = os.listdir(json_folder)
self.json_files = []
for _file in json_files:
if 'manual.json' in _file:
self.json_files.append(_file)
self.video_folder = video_folder
self.bs = bs
self.num_select_frames = select_frames
def __len__(self):
return len(self.json_files) // self.bs
def _get_data(self, idx):
other_infos = {}
json_name = self.json_files[idx]
json_path = os.path.join(self.json_folder, json_name)
with open(json_path, 'r') as f:
data = json.load(f)
other_infos['video_id'] = data['video_id']
video_path = os.path.join(self.video_folder, '{}.mp4'.format(data['video_id']))
frames = get_video_frames(video_path)
masklents = decode_masklet(data['masklet'])
frames = frames[::4]
assert len(frames) == len(masklents)
# frames [np.array(h, w, 3), ...]
# masklents [np.array(h, w, n)]
n_objs = masklents[0].shape[-1]
objects_images = []
for i in range(n_objs):
object_masklents = [_item[:, :, i] for _item in masklents]
select_frame_idxs = self.select_frames(object_masklents, nums=self.num_select_frames)
object_frames = [copy.deepcopy(frames[_idx]) for _idx in select_frame_idxs]
object_masks = [copy.deepcopy(object_masklents[_idx]) for _idx in select_frame_idxs]
object_highlighted_images = self.highlight_object(object_frames, object_masks)
# _folder = os.path.join('./work_dirs/sam2_obj_images', 'obj_{}'.format(i))
# os.mkdir(_folder)
# for j, _save_iamge in enumerate(object_highlighted_images):
# _save_iamge.save(os.path.join(_folder, f'{j}.png'))
question = self.get_question(len(object_highlighted_images))
self._save_drawed_contours(object_highlighted_images,
video_id=other_infos['video_id'],
obj_id=i,
)
objects_images.append({'images': object_highlighted_images, 'text_prompt': question})
return objects_images, other_infos
def _save_drawed_contours(self, images, video_id, obj_id):
for frame_id, image in enumerate(images):
frame_name = f'{video_id}_obj{obj_id}_frame{frame_id}.png'
image.save(os.path.join('/mnt/bn/xiangtai-training-data/project/xiangtai-windows/tt_vlm/work_dirs/object_contour_demos/', frame_name))
return
def get_question(self, num_objs):
ret = ''
for i in range(num_objs):
ret += f'Frame-{i+1}: {IMAGE_TOKEN}\n'
ret += 'Here are several consecutive frames from a video. We have highlighted an object with yellow edges, meaning the object highlighted by the yellow edges in the video is the same object. We need you to provide some discriminative descriptions about this object, which can help us easily distinguish it from other similar objects in the image. The discriminative descriptions should include but are not limited to its category, color, shape, position in the image, state, purpose, properties, and its relationship with surrounding objects.\n'
# ret += 'Please provide a detailed description of the object highlighted by the yellow contour, including its color, shape, position in the image, state, purpose, properties, and its relationship with surrounding objects.'
ret += 'Please give the discriminative descriptions about the object.'
return ret
def highlight_object(self, object_frames, object_masks):
ret = []
for frame, mask in zip(object_frames, object_masks):
image = add_edge_color(frame, mask)
ret.append(image)
return ret
def select_frames(self, object_masklents, nums=3):
areas = np.array([np.sum(mask) for mask in object_masklents])
frame_indexes = np.arange(0, len(object_masklents))
sort_idxs = np.argsort(areas)[::-1]
frame_indexes = frame_indexes[sort_idxs][:nums].tolist()
frame_indexes.sort()
return frame_indexes
def __getitem__(self, idx):
start = idx * self.bs
end = start + self.bs
data_dicts = []
for _idx in range(start, end):
objects_images, other_infos = self._get_data(_idx)
for i, object_dict in enumerate(objects_images):
object_dict.update(other_infos)
object_dict.update({'obj_id': i})
data_dicts.append(object_dict)
return {'data_dicts': data_dicts, 'image_paths': None, 'type': 'sam2'}
@master_only
def evaluate(self, **kwargs):
return {'Acc': 0}
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[:, :, ::-1])
frame_id += 1
cap.release()
return frames
def images_to_video(frames, video_name, fps=6):
height, width, layers = frames[0].shape
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video = cv2.VideoWriter(video_name, fourcc, fps, (width, height))
for frame in frames:
video.write(frame[:, :, ::-1])
# cv2.destroyAllWindows()
video.release()
return
def decode_masklet(masklet):
masks = []
for _rle in masklet:
mask = maskUtils.decode(_rle)
masks.append(mask)
return masks
def draw_mask(image, mask):
obj_mask = mask * 255
obj_mask = np.stack([obj_mask * 1, obj_mask * 0, obj_mask * 0], axis=2)
obj_mask = obj_mask * 0.5 + copy.deepcopy(image) * 0.5
obj_mask = obj_mask.astype(np.uint8)
return obj_mask
def add_mask2images(frames, masklets):
show_videos = []
for i_frames, (frame, masks) in enumerate(zip(frames, masklets)):
if i_frames == 0:
n_obj = masks.shape[-1]
for i_obj in range(n_obj):
show_videos.append([])
n_obj = masks.shape[-1]
for i_obj in range(n_obj):
show_videos[i_obj].append(draw_mask(copy.deepcopy(frame), masks[:, :, i_obj]))
return show_videos
def add_edge_color(image, mask, edge_color=(255, 255, 0), thickness=3):
mask = mask.astype(np.uint8)
contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
tuple_contours = tuple([np.array(contour) for contour in contours])
cv2.drawContours(image, tuple_contours, -1, color=edge_color, thickness=thickness)
image = image.astype(np.uint8)
image = Image.fromarray(image)
return image |