# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: tsat
#     language: python
#     name: tsat
# ---

import torch, torchvision
from torchvision.datasets.video_utils import VideoClips
import os, json
import numpy as np
from torch.utils.data import Dataset, DataLoader
import ffmpeg
from collections import Counter as C


class prober_class(Dataset):
    def __init__(self, paths):
        super(prober_class, self).__init__()
        self.paths = paths

    def __len__(self):
        return len(self.paths)
    
    def __getitem__(self, idx):
        k,p = self.paths[idx]
        assert os.path.exists(p)
        try:
            res = ffmpeg.probe(p)
        except:
            res = "Error"
        
        return [res, k]


# +
basepath = '/data/datasets/kinetics700_2020/'

def read_dataset_details(label_path, batch_size, vid_path):
    print("hello")
    dataset = json.load(open(basepath + label_path))
    print(list(dataset.items())[0], len(dataset))
    collate = lambda x: x
    loader = DataLoader(prober_class([(k, basepath + vid_path + k + '.mp4') for k, v in dataset.items()]),
                         batch_size=batch_size,
                         num_workers=batch_size,
                         collate_fn=collate)
    print("Example paths:", loader.dataset.paths[:3])
    probes = []
    for i,batch in enumerate(loader):
        if i % 50 == 0:
            print(i*batch_size,'/',len(dataset))
        for b in batch:
            probes.append(b)
    print(C([type(p[0]) for p in probes]))
    return probes


# -

def filter_dataset(label_path, probes):
    dataset = json.load(open(basepath + label_path))
    print(list(dataset.items())[0], len(dataset))
    
    new_dataset = {}
    error_remove_ids = []
    stream_err_ids = []

    for p, k in probes:
        if p == "Error":
            error_remove_ids.append(k)
        elif len(p['streams']) != 2:
            stream_err_ids.append(k)
        else:
            dataset[k]['nb_frames'] = p['streams'][0]['nb_frames']
            dataset[k]['hw'] = (p['streams'][0]['height'], p['streams'][0]['width'])
            dataset[k]['true_duration'] = p['streams'][0]['duration']
            new_dataset[k] = dataset[k]
            
    return new_dataset, error_remove_ids, stream_err_ids


probess = read_dataset_details('full/train.json', 50, 'train_cache_short_fixed/')

len(probess)

probess[1][0]['streams']













new_dataset, error_remove_ids, stream_err_ids = filter_dataset('full/train.json', probess)

len(stream_err_ids)

# +
#json.dump(new_dataset, open(basepath + 'labels/val.json', 'w+'), indent=2)
# -

probess = read_dataset_details('labels/train_full.json', 60, 'train/')

json.dump(probess, open("tmp.json", 'w+'))

new_dataset, error_remove_ids, stream_err_ids = filter_dataset('labels/train_full.json', probess)

len(new_dataset)

len(error_remove_ids)

json.dump(error_remove_ids, open("corrupt_missing_train.json", 'w+'))

json.dump(stream_err_ids, open("stream_err_train.json", 'w+'))

json.dump(['xsGtbUnp9tw', '141fJ89Ed2k', 'q-S0pBZDhZU', 'wN8qYmPv5yk',
           'uz6rjbw0ZA0', 'kUVnT7Ld80M', 'XY5FDVay5_A', 'bOU2oGVBM_o',
           'fg2BS7H_dAU', 'KTCQpjUrCe8', 'erh2ngRZxs0', '7hIAtSLdAUo',
           'h_5SZwWFg1c', 'z8qEtdr1ZuU', 'VYZPozZ5Eig', '6iuD3pSgBcw',
           'NPNP-7B9P3M', 'yhA_TTKetyM', 'YjUcA9zOp5g', 'hKqCUWTQQxU'
          ],
         open("download_err_corrupt.json", 'w+')
        )

len(new_dataset)

json.dump(new_dataset, open(basepath + 'labels/train.json', 'w+'), indent=2)

# Create mini kinetics train
import random
label_class = list(set(m['annotations']['label'] for m in new_dataset.values()))
assert len(label_class) == 700
classwise_keys = {l:[] for l in label_class}
for k, d in new_dataset.items():
    classwise_keys[d['annotations']['label']].append(k)


classwise_ixs['sword swallowing']

mini_train = dict(random.sample(list(new_dataset.items()), 200000))

dict(sorted(C(m['annotations']['label'] for m in mini_train.values()).items(),key=lambda x: x[1]))

label_to_ix = {k:i for i, k in enumerate(label_class)}

json.dump(label_to_ix, open(basepath + 'mini/class_to_index.json', 'w+'), indent=2)

for k in mini_train.keys():
    l = mini_train[k]['annotations']['label']
    mini_train[k]['annotations']['label'] = (l, label_to_ix[l])

list(mini_train.items())[1], len(mini_train)

json.dump(mini_train, open(basepath + 'mini/train.json', 'w+'), indent=2)

# +
mini_val = json.load(open(basepath + 'labels/val.json'))

for k in mini_val.keys():
    l = mini_val[k]['annotations']['label']
    mini_val[k]['annotations']['label'] = (l, label_to_ix[l])
# -

list(mini_val.items())[1], len(mini_val)

json.dump(mini_val, open(basepath + 'mini/val.json', 'w+'), indent=2)



StarcoderdataJupyterScriptsDedupFiltered