import numpy as np
from datasets import load_dataset
dataset = load_dataset('WenhaoWang/VidProM')

import os
import pandas as pd
import h5py
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.io import read_video
import numpy as np

class VidProMDataset(Dataset):
    def __init__(self, csv_file, hdf5_file, video_dirs, transform=None):
        
        self.metadata = pd.read_csv(csv_file)
        self.video_dirs = video_dirs
        self.transform = transform
        self.nsfw_names = ['toxicity','obscene','identity_attack','insult','threat','sexual_explicit']

        self.hdf5_file =  h5py.File(hdf5_file, 'r')
        self.hdf5_uuid = np.array(self.hdf5_file["uuid"][:], dtype=object).astype(str).tolist()
        self.hdf5_embed = np.array(self.hdf5_file['embeddings'])
        
    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, idx):
        
        video_info = self.metadata.iloc[idx]
        video_id = video_info['uuid']
        prompt = video_info['prompt']
        time = video_info['time']
        nsfw_scores = torch.tensor(list(video_info[self.nsfw_names]))
        
        embed = torch.tensor(self.hdf5_embed[self.hdf5_uuid.index(video_id)])
        video_path = self._find_video_path(video_id)
        video_frames, _, _ = read_video(video_path, pts_unit='sec')
        
        if self.transform:
            video_frames = self.transform(video_frames)

        return {
            'video_id': video_id,
            'video_frames': video_frames,
            'embed': embed,
            'prompt': prompt,
            'time': time,
            'nsfw_scores': nsfw_scores
        }

    def _find_video_path(self, video_id):
        for video_dir in self.video_dirs:
            video_file = os.path.join(video_dir, video_dir.split('_')[0] + f"-{video_id}.mp4")
            if os.path.exists(video_file):
                return video_file
        raise FileNotFoundError(f"Video {video_id}.mp4 not found in any of the directories.")

    def __del__(self):
        self.hdf5_file.close()

csv_file = 'VidProM_unique_example.csv'
hdf5_file = 'VidProM_embed_example.hdf5'
video_dirs = ['t2vz_videos_example', 'pika_videos_example', 'vc2_videos_example', 'ms_videos_example']
dataset = VidProMDataset(csv_file, hdf5_file, video_dirs)
dataloader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=0)

import os
import io
import av
import pandas as pd
import h5py
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import torch
import webdataset as wds

tar_file_path = 't2vz_videos_example.tar' # we use t2vz_videos_example.tar for example
csv_file = 'VidProM_unique_example.csv'
hdf5_file = 'VidProM_embed_example.hdf5'
dataset = wds.WebDataset(tar_file_path)
metadata = pd.read_csv(csv_file)
hdf5_file = h5py.File(hdf5_file, 'r')
hdf5_uuid = np.array(hdf5_file["uuid"][:], dtype=object).astype(str).tolist()
hdf5_embed = np.array(hdf5_file['embeddings'])

for sample in dataset:
    #obtain tensor of a video
    binary_data = sample['mp4']
    container = av.open(io.BytesIO(binary_data))
    transform = transforms.ToTensor()
    frames = []
    for frame in container.decode(video=0):
        img = frame.to_image()  
        img_tensor = transform(img) 
        frames.append(img_tensor)  
    video_tensor = torch.stack(frames)
    
    #obtain uuid of a video
    uuid = '-'.join(sample['__key__'].split('/')[-1].split('-')[1:])
    
    #obtain the prompt
    prompt = list(metadata[metadata['uuid']==uuid].iloc[:, 1])[0]
    
    #obtain the time
    time = list(metadata[metadata['uuid']==uuid].iloc[:, 2])[0]
    
    #obtain the nsfw_scores
    nsfw_scores = list(metadata[metadata['uuid']==uuid].iloc[0, 3:])
    
    #obtain the prompt embedding
    embed = torch.tensor(hdf5_embed[hdf5_uuid.index(uuid)])

import pandas
df = pd.read_csv("VidProM_unique.csv")

import numpy as np
import h5py
def read_descriptors(filename):
    hh = h5py.File(filename, "r")
    descs = np.array(hh["embeddings"])
    names = np.array(hh["uuid"][:], dtype=object).astype(str).tolist()
    return names, descs

uuid, features = read_descriptors('VidProM_embed.hdf5')