Spaces:
Running
on
Zero
Running
on
Zero
| from typing import Dict, List | |
| import random | |
| from PIL import Image, ImageDraw, ImageFont | |
| from .utils import sample_video | |
| class VideoPermutationParser: | |
| def __init__( | |
| self, | |
| n_frames=8, | |
| is_training=True, | |
| frame_nums = list(range(8, 25)), | |
| video_sampling_strategy={}, | |
| ): | |
| self.n_frames = n_frames | |
| self.is_training = is_training | |
| self.frame_nums = frame_nums | |
| self.video_sampling_strategy = video_sampling_strategy | |
| # fmt: off | |
| self.data_temp = { | |
| "text": [{ | |
| "prompt": "<video>", | |
| "response": "" | |
| }], | |
| "video": [{ | |
| "video_file": { | |
| "yg": "/mnt/bn/videonasyg/videos/webvid_10M_download/011851_011900/1047443473.mp4", | |
| "lq": "/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] | |
| }], | |
| } | |
| # fmt: on | |
| def check_format(self, data_dict: Dict): | |
| pass | |
| # for k in self.data_temp.keys(): | |
| # assert k in data_dict | |
| def transform(self, data_dict: Dict, image_processing_config: Dict = None) -> Dict: | |
| self.check_format(data_dict) | |
| frames = self.load_video_item(data_dict['video'][0]) | |
| # frames = self.add_text_to_frames(frames) # for debug | |
| idxs = list(range(1, len(frames) + 1)) | |
| random.shuffle(idxs) | |
| prefix_len = int(3/8*len(idxs)) | |
| shuffled_frames = [frames[i-1] for i in idxs] | |
| prompt = f'Output the correct chronological order of scrambled video frames. The order of the first {prefix_len} ones are:\n' | |
| prompt += '\n'.join([str(i) for i in idxs[: prefix_len]]) + '\nOutput the order of the following frames:' | |
| response = '\n'.join([str(i) for i in idxs[prefix_len: ]]) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "video", "video": shuffled_frames}, | |
| {"type": "text", "text": prompt}, | |
| ] | |
| }, | |
| { | |
| "role": "assistant", | |
| "content": [ | |
| {"type": "text", "text": response} | |
| ] | |
| } | |
| ] | |
| return messages | |
| def load_video_item(self, video_item) -> List[Image.Image]: | |
| """ | |
| video_item: | |
| {"video_file": "/path/to/video", "n_frames": 8} | |
| {"video_file": "/path/to/video", "frame_indices": [0, 1, 2], "n_frames": 3} | |
| {"video_file": "/path/to/video", "start_frame": 0, "end_frame": 100, "n_frames": 8} | |
| {"video_file": "/path/to/video", "time_indices": [0, 1, 2], "n_frames": 3} | |
| {"video_file": "/path/to/video", "start_time": 0, "end_time": 100, "n_frames": 8} | |
| {"image_file": ["/path/to/image"], "frame_indices": [0, 1, 2], "n_frames": 3} | |
| """ | |
| # check format | |
| if ("image_file" not in video_item) and ("video_file" not in video_item): | |
| raise KeyError(f"Key 'image_file' or 'video_file' not found in video_item") | |
| video_path = video_item.get('video_file', video_item.get('image_file')) | |
| n_frames = video_item.get('n_frames', None) | |
| frame_indices = video_item.get('frame_indices', None) | |
| start_frame = video_item.get('start_frame', None) | |
| end_frame = video_item.get('end_frame', None) | |
| time_indices = video_item.get('time_indices', None) | |
| start_time = video_item.get('start_time', None) | |
| end_time = video_item.get('end_time', None) | |
| mask_boxes = video_item.get('mask_boxes', None) | |
| n_frames = random.choice(self.frame_nums) | |
| n = self.video_sampling_strategy.get('force_frames_n_divisible', 1) | |
| if n > 1 and n_frames % n != 0: | |
| n_frames += n - n_frames % n | |
| frames, frame_indices = sample_video( | |
| video_path=video_path, | |
| frame_indices=frame_indices, | |
| start_frame=start_frame, | |
| end_frame=end_frame, | |
| n_frames=n_frames, | |
| time_indices=time_indices, | |
| start_time=start_time, | |
| end_time=end_time, | |
| mask_boxes=mask_boxes, | |
| is_training=self.is_training, | |
| video_sampling_strategy=self.video_sampling_strategy, | |
| return_frame_ids=True, | |
| ) | |
| return frames | |
| def add_text_to_frames(self, frames: List[Image.Image]): | |
| new_frames = [] | |
| for i, image in enumerate(frames): | |
| draw = ImageDraw.Draw(image) | |
| font = ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf', 100) | |
| text_position = (50, 50) | |
| text_content = f'{i+1}' | |
| text_color = (255, 0, 0) | |
| draw.text(text_position, text_content, font=font, fill=text_color) | |
| new_frames.append(image) | |
| return new_frames | |