diff --git a/model/dataset.py b/model/dataset.py
index dae04a7..84ccaf9 100644
--- a/model/dataset.py
+++ b/model/dataset.py
@@ -4,10 +4,48 @@ 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):
@@ -84,3 +122,80 @@ class VLMDataset(Dataset):
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('assistant\n', add_special_tokens=False).input_ids
+ self.eos_id = tokenizer('\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('