202 lines
7.7 KiB
Python
202 lines
7.7 KiB
Python
import json
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset, DataLoader
|
|
import torch
|
|
from .model_vlm import MiniMindVLM
|
|
import os
|
|
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode
|
|
import cv2
|
|
import numpy as np
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
|
|
def video2image(video_path, num_frames=8, size=224):
|
|
def preprocess(size, n_px):
|
|
return Compose([
|
|
Resize(size, interpolation=InterpolationMode.BICUBIC),
|
|
CenterCrop(size),
|
|
lambda image: image.convert("RGB"),
|
|
ToTensor(),
|
|
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
|
|
])(n_px)
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
|
|
frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
|
|
|
if fps < 1 or frameCount < 1:
|
|
images = np.zeros([3, size, size], dtype=np.float32)
|
|
print("ERROR: problem reading video file: ", video_path)
|
|
else:
|
|
frames_idx = np.sort(np.random.choice(frameCount, min(num_frames, frameCount), replace=False))
|
|
|
|
images = np.zeros([len(frames_idx), 3, size, size], dtype=np.float32)
|
|
|
|
for i, idx in enumerate(frames_idx):
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
continue
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
images[i,:,:,:] = preprocess(size, Image.fromarray(frame).convert("RGB"))
|
|
|
|
cap.release()
|
|
video_frames = torch.tensor(images)
|
|
return video_frames
|
|
|
|
class VLMDataset(Dataset):
|
|
def __init__(self, jsonl_path, images_path, tokenizer, preprocess=None, max_length=512,
|
|
image_special_token='@' * 196):
|
|
|
|
super().__init__()
|
|
self.samples = self.load_data(jsonl_path)
|
|
self.images_path = images_path
|
|
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
self.preprocess = preprocess
|
|
self.image_token = image_special_token
|
|
self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
|
|
self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def load_data(self, path):
|
|
samples = []
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
for line_num, line in enumerate(f, 1):
|
|
data = json.loads(line.strip())
|
|
samples.append(data)
|
|
return samples
|
|
|
|
def _create_chat_prompt(self, conversations):
|
|
messages = []
|
|
for i, turn in enumerate(conversations):
|
|
role = 'user' if i % 2 == 0 else 'assistant'
|
|
messages.append({"role": role, "content": turn['content'].replace('<image>', self.image_token)})
|
|
return self.tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False
|
|
)
|
|
|
|
def _generate_loss_mask(self, input_ids):
|
|
loss_mask = [0] * len(input_ids)
|
|
i = 0
|
|
while i < len(input_ids):
|
|
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
|
start = i + len(self.bos_id)
|
|
end = start
|
|
while end < len(input_ids):
|
|
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
|
break
|
|
end += 1
|
|
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
|
|
loss_mask[j] = 1
|
|
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
|
else:
|
|
i += 1
|
|
return loss_mask
|
|
|
|
def __getitem__(self, index: int):
|
|
sample = self.samples[index]
|
|
image_paths = sample['image']
|
|
prompt = self._create_chat_prompt(sample['conversations'])
|
|
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
|
|
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
|
|
loss_mask = self._generate_loss_mask(input_ids)
|
|
|
|
X = torch.tensor(input_ids[:-1], dtype=torch.long)
|
|
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
|
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
|
|
|
|
image_tensors = []
|
|
for image_name in image_paths.split(','):
|
|
image_name = image_name.strip()
|
|
image = Image.open(f'{self.images_path}/{image_name}')
|
|
image_tensor = MiniMindVLM.image2tensor(image, self.preprocess)
|
|
image_tensors.append(image_tensor)
|
|
image_tensors = torch.stack(image_tensors, dim=0)
|
|
|
|
return X, Y, loss_mask, image_tensors
|
|
|
|
|
|
class VideoLMDataset(Dataset):
|
|
def __init__(self, jsonl_path, videos_path, tokenizer, max_length=512,
|
|
video_special_token='@' * 196):
|
|
|
|
super().__init__()
|
|
self.samples = self.load_data(jsonl_path)
|
|
self.videos_path = videos_path
|
|
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
self.video_token = video_special_token
|
|
self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
|
|
self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def load_data(self, path):
|
|
samples = []
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
for line_num, line in enumerate(f, 1):
|
|
data = json.loads(line.strip())
|
|
samples.append(data)
|
|
return samples
|
|
|
|
def _create_chat_prompt(self, conversations):
|
|
messages = []
|
|
for i, turn in enumerate(conversations):
|
|
role = 'user' if i % 2 == 0 else 'assistant'
|
|
messages.append({"role": role, "content": turn['content'].replace('<video>', self.video_token)})
|
|
return self.tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False
|
|
)
|
|
|
|
def _generate_loss_mask(self, input_ids):
|
|
loss_mask = [0] * len(input_ids)
|
|
i = 0
|
|
while i < len(input_ids):
|
|
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
|
start = i + len(self.bos_id)
|
|
end = start
|
|
while end < len(input_ids):
|
|
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
|
break
|
|
end += 1
|
|
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
|
|
loss_mask[j] = 1
|
|
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
|
else:
|
|
i += 1
|
|
return loss_mask
|
|
|
|
def __getitem__(self, index: int):
|
|
sample = self.samples[index]
|
|
video_paths = sample['video']
|
|
prompt = self._create_chat_prompt(sample['conversations'])
|
|
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
|
|
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
|
|
loss_mask = self._generate_loss_mask(input_ids)
|
|
|
|
X = torch.tensor(input_ids[:-1], dtype=torch.long)
|
|
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
|
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
|
|
|
|
video_tensors = []
|
|
for video_name in video_paths.split(','):
|
|
video_name = video_name.strip()
|
|
video_path = f'{self.videos_path}/{video_name}'
|
|
video_tensor = video2image(video_path)
|
|
video_tensors.append(video_tensor)
|
|
video_tensors = torch.stack(video_tensors, dim=0)
|
|
|
|
return X, Y, loss_mask, video_tensors
|