Spaces:
Runtime error
Runtime error
| from typing import Dict, List | |
| import random | |
| import re | |
| from PIL import Image | |
| from .utils import sample_video, read_image | |
| class MultiImagesParser: | |
| def __init__( | |
| self, | |
| n_frames=8, | |
| is_training=True, | |
| ): | |
| self.n_frames = n_frames | |
| self.is_training = is_training | |
| # fmt: off | |
| self.data_temp = { | |
| "text": [ | |
| [{ | |
| "prompt": "Describe the image in short.", | |
| "response": "A rollerblader rides high in a full pipe while others watch" | |
| }], | |
| [{ | |
| "prompt": "Describe the image in short.", | |
| "response": "A woman in winter clothes is on the sidewalk with a phone." | |
| }] | |
| ], | |
| "image": [ | |
| { | |
| "image_file": "/mnt/bn/videonaslq/images/flickr30k/images/3371533654.jpg" | |
| }, | |
| { | |
| "image_file": "/mnt/bn/videonaslq/images/coco/train2014/COCO_train2014_000000177950.jpg" | |
| }, | |
| { | |
| "video_file": "/mnt/bn/llmdatalq/jiangnan/video_generation/webvid_10M_download/20230609/videos/011851_011900/1047443473.mp4", | |
| "frame_indices": [0, 85, 171, 256, 342, 427, 513, 598] | |
| } | |
| ], | |
| "dataset": "coco", | |
| "task": "multi_images", | |
| "image_processing_config": {}, | |
| } | |
| # fmt: on | |
| def check_format(self, data_dict: Dict, image_processing_config: Dict): | |
| assert data_dict['dataset'] in ['coco', 'sharegpt4v_cap100k', 'sharegpt4v_mix665k', 'webvid', 'movie'], data_dict | |
| # 目前多图数据应该没有包含坐标的数据吧 | |
| if image_processing_config.get('has_coordinates', False): | |
| raise ValueError(f'do_crop and has_coordinates cannot be True at the same time in MultiImagesParser!') | |
| # 检查是否能匹配到坐标 | |
| texts = data_dict['text'] | |
| for text in texts: | |
| match = re.search(r'\[(\d+(\.\d+)?,\s*)+\d+(\.\d+)?\]', text['prompt'] + text['response']) | |
| if match: | |
| print(f'[Warning] 疑似检测到包含坐标的数据：{data_dict}') | |
| def transform(self, data_dict: Dict, image_processing_config: Dict = None) -> Dict: | |
| self.check_format(data_dict, image_processing_config) | |
| # shuffle | |
| texts = data_dict['text'] | |
| images = data_dict['image'] | |
| images = self.load_images(images) | |
| idxs = list(range(len(texts))) | |
| random.shuffle(idxs) | |
| texts = [texts[i] for i in idxs] | |
| images = [images[i] for i in idxs] | |
| # sample n_frames | |
| if isinstance(self.n_frames, int): | |
| n_frames = random.choice(list(range(1, self.n_frames + 1))) | |
| else: | |
| n_frames = random.choice(self.n_frames) | |
| texts = texts[: n_frames] | |
| images = images[: n_frames] | |
| dataset = data_dict['dataset'] | |
| if dataset in ['coco', 'sharegpt4v_cap100k', 'webvid', 'movie']: | |
| prompt, response = self.transform_for_caption_task(texts, dataset, images) | |
| else: | |
| prompt, response = self.transform_for_qa_task(texts, dataset, images) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| *[{"type": "image", "image": img} for img in images], | |
| {"type": "text", "text": prompt}, | |
| ] | |
| }, | |
| { | |
| "role": "assistant", | |
| "content": [ | |
| {"type": "text", "text": response} | |
| ] | |
| } | |
| ] | |
| return messages | |
| def transform_for_caption_task(self, texts, dataset, images): | |
| idx = random.choice(list(range(len(texts)))) | |
| if dataset == 'coco': | |
| if len(texts) == 1: | |
| prompt = 'Describe the image in short.' | |
| else: | |
| prompt = f'Describe the images starting from frame {idx + 1} in short in order.' | |
| elif dataset == 'sharegpt4v_cap100k': | |
| if len(texts) == 1: | |
| prompt = 'Describe the image in detail.' | |
| else: | |
| prompt = f'Describe the images starting from frame {idx + 1} in detail in order.' | |
| else: | |
| if len(texts) == 1: | |
| prompt = 'Describe the image.' | |
| else: | |
| prompt = f'Describe the images starting from frame {idx + 1} in order.' | |
| response = '' | |
| for i, text in enumerate(texts): | |
| if i < idx: | |
| continue | |
| if not isinstance(text, dict): | |
| text = random.choice(text) | |
| resp = text['response'] | |
| response += f'{resp}\n' | |
| return prompt, response | |
| def transform_for_qa_task(self, texts, dataset, images): | |
| prompt, response = '', '' | |
| for i, text in enumerate(texts): | |
| if not isinstance(text, dict): | |
| text = random.choice(text) | |
| if len(texts) > 1: | |
| prompt += f'Question for frame {i+1}:\n' + text['prompt'] + '\n' | |
| response += f'Answer to question of frame {i+1}:\n' + text['response'] + '\n' | |
| else: | |
| prompt += text['prompt'] + '\n' | |
| response += text['response'] + '\n' | |
| return prompt, response | |
| def load_images(self, image_items: List[Dict]) -> List[Image.Image]: | |
| """ | |
| image_items: List[Dict]. each item like: | |
| {"video_file": "path/to/video", "frame_indices": [1]} | |
| or | |
| {"image_file": "path/to/image"} | |
| """ | |
| if image_items is None: | |
| raise ValueError(f'image_items is None!') | |
| if isinstance(image_items, dict): | |
| image_items = [image_items] | |
| images = [] | |
| for image_item in image_items: | |
| if 'video_file' in image_item: | |
| file_key = 'video_file' | |
| elif 'image_file' in image_item: | |
| file_key = 'image_file' | |
| else: | |
| raise KeyError(f'video_file or image_file not in {image_item}') | |
| file_path = image_item[file_key] | |
| if file_key == 'video_file': | |
| frame_indices = image_item.get('frame_indices', None) | |
| if frame_indices is None: | |
| raise ValueError(f'read 0 frame: {image_item}') | |
| if isinstance(frame_indices, int): | |
| frame_indices = [frame_indices] | |
| frames = sample_video(file_path, frame_indices = frame_indices) | |
| images.extend(frames) | |
| else: | |
| if isinstance(file_path, str): | |
| file_path = [file_path] | |
| images.extend([read_image(f) for f in file_path]) | |
| return images | |
| if __name__ == '__main__': | |
| # python3 -m xenon_generation.data.custom_data_parsers.multi_images_parser | |
| from tqdm import tqdm | |
| from tools.rw_utils import read_jsonlines | |
| lines = read_jsonlines('/mnt/bn/videonaslq/VideoCaption/datasets_1009/sharegpt4v_cap100k/part_36.jsonl') | |
| lines = lines[:10] | |
| parser = MultiImagesParser(n_frames=8) | |
| for i, l in tqdm(enumerate(lines)): | |
| l_image_processing_config = l.get('image_processing_config', {}) | |
| messages = parser.transform(l, l_image_processing_config) | |
| print(messages) |