chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,5 @@
|
||||
dataset:
|
||||
bert_name: bert-base-uncased
|
||||
caption_pkl_path: data/how2/raw_caption_dedup.pkl
|
||||
use_fast: true
|
||||
target_dir: data/feat/feat_how2_s3d_shard_small
|
||||
@@ -0,0 +1,106 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import pickle
|
||||
import os
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from mmpt.processors import PKLJSONStrTextProcessor
|
||||
from mmpt.utils import ShardedTensor, recursive_config
|
||||
|
||||
|
||||
class TokenizerDataset(Dataset):
|
||||
def __init__(self, config):
|
||||
self.text_processor = PKLJSONStrTextProcessor(config)
|
||||
self.video_ids = list(self.text_processor.data.keys())
|
||||
|
||||
def __getitem__(self, idx):
|
||||
video_id = self.video_ids[idx]
|
||||
return video_id, self.text_processor(video_id)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.video_ids)
|
||||
|
||||
|
||||
def numpify(shard_idx, video_ids, captions, target_dir, split, prefix, max_cap_len=32):
|
||||
startends = []
|
||||
caps_ids = []
|
||||
for video_id in video_ids:
|
||||
caption = captions[video_id]
|
||||
startend = []
|
||||
cap_ids = []
|
||||
for start, end, cap in zip(
|
||||
caption["start"], caption["end"], caption["cap"]):
|
||||
startend.append(np.array([start, end]).astype("float32"))
|
||||
cap_id = np.full((max_cap_len,), -1, dtype=np.int32)
|
||||
cap = cap[:max_cap_len]
|
||||
cap_id[:len(cap)] = cap
|
||||
cap_ids.append(cap_id)
|
||||
startends.append(np.stack(startend))
|
||||
caps_ids.append(np.stack(cap_ids))
|
||||
|
||||
startends = ShardedTensor.from_list(startends)
|
||||
target_path = os.path.join(
|
||||
target_dir,
|
||||
prefix + split + "_" + str(shard_idx)
|
||||
)
|
||||
print("save to", target_path)
|
||||
startends.save(target_path + ".startends")
|
||||
caps_ids = ShardedTensor.from_list(caps_ids)
|
||||
caps_ids.save(target_path + ".caps_ids")
|
||||
|
||||
|
||||
def sharding(config, out_file):
|
||||
with open(out_file, "rb") as fr:
|
||||
captions = pickle.load(fr)
|
||||
target_dir = config.target_dir
|
||||
prefix = os.path.basename(
|
||||
os.path.splitext(config.caption_pkl_path)[0]
|
||||
) + "." + config.bert_name + "."
|
||||
for split in ["train", "val"]:
|
||||
target_path = os.path.join(target_dir, split + "_meta")
|
||||
with open(target_path + ".pkl", "rb") as fr:
|
||||
meta = pickle.load(fr)
|
||||
print("load meta", target_path, len(meta))
|
||||
for shard_id in meta:
|
||||
numpify(
|
||||
shard_id, meta[shard_id], captions,
|
||||
target_dir, split, prefix
|
||||
)
|
||||
|
||||
|
||||
def tokenize(config, out_file):
|
||||
def collator(samples):
|
||||
return samples
|
||||
dataset = TokenizerDataset(config)
|
||||
data = {}
|
||||
for idx, batch in enumerate(
|
||||
DataLoader(dataset, collate_fn=collator, num_workers=16)):
|
||||
for video_id, caption in batch:
|
||||
data[video_id] = caption
|
||||
if idx % 5000 == 0:
|
||||
print(idx)
|
||||
with open(out_file, "wb") as fw:
|
||||
pickle.dump(data, fw, pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
|
||||
def main(args):
|
||||
config = recursive_config(args.config).dataset
|
||||
|
||||
out_file = os.path.splitext(config.caption_pkl_path)[0] \
|
||||
+ "." + config.bert_name + ".pkl"
|
||||
if not os.path.isfile(out_file):
|
||||
tokenize(config, out_file)
|
||||
sharding(config, out_file)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="pretokenize (raw_)caption.json into pkl.")
|
||||
parser.add_argument('config', type=str)
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -0,0 +1,157 @@
|
||||
# Copyright Howto100M authors.
|
||||
# Copyright (c) Facebook, Inc. All Rights Reserved
|
||||
|
||||
import torch as th
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
import numpy as np
|
||||
import argparse
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
from model import get_model
|
||||
from preprocessing import Preprocessing
|
||||
from random_sequence_shuffler import RandomSequenceSampler
|
||||
|
||||
from tqdm import tqdm
|
||||
from pathbuilder import PathBuilder
|
||||
from videoreader import VideoLoader
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description='Easy video feature extractor')
|
||||
|
||||
parser.add_argument('--vdir', type=str)
|
||||
parser.add_argument('--fdir', type=str)
|
||||
parser.add_argument('--hflip', type=int, default=0)
|
||||
|
||||
parser.add_argument('--batch_size', type=int, default=64,
|
||||
help='batch size')
|
||||
parser.add_argument('--type', type=str, default='2d',
|
||||
help='CNN type')
|
||||
parser.add_argument('--half_precision', type=int, default=0,
|
||||
help='output half precision float')
|
||||
parser.add_argument('--num_decoding_thread', type=int, default=4,
|
||||
help='Num parallel thread for video decoding')
|
||||
parser.add_argument('--l2_normalize', type=int, default=1,
|
||||
help='l2 normalize feature')
|
||||
parser.add_argument('--resnext101_model_path', type=str, default='model/resnext101.pth',
|
||||
help='Resnext model path')
|
||||
parser.add_argument('--vmz_model_path', type=str, default='model/r2plus1d_34_clip8_ig65m_from_scratch-9bae36ae.pth',
|
||||
help='vmz model path')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
# TODO: refactor all args into config. (current code is from different people.)
|
||||
CONFIGS = {
|
||||
"2d": {
|
||||
"fps": 1,
|
||||
"size": 224,
|
||||
"centercrop": False,
|
||||
"shards": 0,
|
||||
},
|
||||
"3d": {
|
||||
"fps": 24,
|
||||
"size": 112,
|
||||
"centercrop": True,
|
||||
"shards": 0,
|
||||
},
|
||||
"s3d": {
|
||||
"fps": 30,
|
||||
"size": 224,
|
||||
"centercrop": True,
|
||||
"shards": 0,
|
||||
},
|
||||
"vmz": {
|
||||
"fps": 24,
|
||||
"size": 112,
|
||||
"centercrop": True,
|
||||
"shards": 0,
|
||||
},
|
||||
"vae": {
|
||||
"fps": 2,
|
||||
"size": 256,
|
||||
"centercrop": True,
|
||||
"shards": 100,
|
||||
}
|
||||
}
|
||||
|
||||
config = CONFIGS[args.type]
|
||||
|
||||
|
||||
video_dirs = args.vdir
|
||||
feature_dir = args.fdir
|
||||
|
||||
video_dict = PathBuilder.build(video_dirs, feature_dir, ".npy", config["shards"])
|
||||
|
||||
dataset = VideoLoader(
|
||||
video_dict=video_dict,
|
||||
framerate=config["fps"],
|
||||
size=config["size"],
|
||||
centercrop=config["centercrop"],
|
||||
hflip=args.hflip
|
||||
)
|
||||
n_dataset = len(dataset)
|
||||
sampler = RandomSequenceSampler(n_dataset, 10)
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=1,
|
||||
shuffle=False,
|
||||
num_workers=args.num_decoding_thread,
|
||||
sampler=sampler if n_dataset > 10 else None,
|
||||
)
|
||||
preprocess = Preprocessing(args.type)
|
||||
model = get_model(args)
|
||||
|
||||
with th.no_grad():
|
||||
for k, data in tqdm(enumerate(loader), total=loader.__len__(), ascii=True):
|
||||
input_file = data['input'][0]
|
||||
output_file = data['output'][0]
|
||||
if len(data['video'].shape) > 3:
|
||||
video = data['video'].squeeze()
|
||||
if len(video.shape) == 4:
|
||||
video = preprocess(video)
|
||||
n_chunk = len(video)
|
||||
if args.type == 'vmz':
|
||||
n_chunk = math.ceil(n_chunk/float(3))
|
||||
features = th.cuda.FloatTensor(n_chunk, 512).fill_(0)
|
||||
elif args.type == 's3d':
|
||||
features = th.cuda.FloatTensor(n_chunk, 512).fill_(0)
|
||||
elif args.type == "vae":
|
||||
features = th.cuda.LongTensor(n_chunk, 1024).fill_(0)
|
||||
else:
|
||||
features = th.cuda.FloatTensor(n_chunk, 2048).fill_(0)
|
||||
n_iter = int(math.ceil(n_chunk / float(args.batch_size)))
|
||||
for i in range(n_iter):
|
||||
factor = 1
|
||||
if args.type == 'vmz':
|
||||
factor = 3
|
||||
min_ind = factor * i * args.batch_size
|
||||
max_ind = factor * (i + 1) * args.batch_size
|
||||
video_batch = video[min_ind:max_ind:factor].cuda()
|
||||
if args.type == '2d':
|
||||
batch_features = model(video_batch) # (51, 487), (51, 512)
|
||||
elif args.type == 's3d':
|
||||
batch_features = model(video_batch)
|
||||
batch_features = batch_features['video_embedding']
|
||||
elif args.type == "vae":
|
||||
# image_code.
|
||||
batch_features = model(video_batch)
|
||||
else:
|
||||
batch_pred, batch_features = model(video_batch) # (51, 487), (51, 512)
|
||||
if args.l2_normalize:
|
||||
batch_features = F.normalize(batch_features, dim=1)
|
||||
features[i*args.batch_size:(i+1)*args.batch_size] = batch_features
|
||||
features = features.cpu().numpy()
|
||||
if args.half_precision:
|
||||
if args.type == "vae":
|
||||
features = features.astype(np.int16)
|
||||
else:
|
||||
features = features.astype('float16')
|
||||
else:
|
||||
if args.type == "vae":
|
||||
features = features.astype(np.int32)
|
||||
else:
|
||||
features = features.astype('float32')
|
||||
np.save(output_file, features)
|
||||
else:
|
||||
print('Video {} error.'.format(input_file))
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
|
||||
|
||||
python scripts/video_feature_extractor/extract.py \
|
||||
--vdir <path_to_video_folder> \
|
||||
--fdir data/feat/feat_how2_s3d \
|
||||
--type=s3d --num_decoding_thread=4 \
|
||||
--batch_size 32 --half_precision 1
|
||||
@@ -0,0 +1,58 @@
|
||||
# Copyright (c) Howto100M authors and Facebook, Inc. All Rights Reserved
|
||||
|
||||
import torch as th
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
class GlobalAvgPool(nn.Module):
|
||||
def __init__(self):
|
||||
super(GlobalAvgPool, self).__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return th.mean(x, dim=[-2, -1])
|
||||
|
||||
|
||||
def get_model(args):
|
||||
assert args.type in ['2d', '3d', 'vmz', 's3d', 'vae']
|
||||
if args.type == '2d':
|
||||
print('Loading 2D-ResNet-152 ...')
|
||||
import torchvision.models as models
|
||||
model = models.resnet152(pretrained=True)
|
||||
model = nn.Sequential(*list(model.children())[:-2], GlobalAvgPool())
|
||||
model = model.cuda()
|
||||
elif args.type == 'vmz':
|
||||
print('Loading VMZ ...')
|
||||
from vmz34 import r2plus1d_34
|
||||
model = r2plus1d_34(pretrained_path=args.vmz_model_path, pretrained_num_classes=487)
|
||||
model = model.cuda()
|
||||
elif args.type == 's3d':
|
||||
# we use one copy of s3d instead of dup another one for feature extraction.
|
||||
from mmpt.processors.models.s3dg import S3D
|
||||
model = S3D('pretrained_models/s3d_dict.npy', 512)
|
||||
model.load_state_dict(th.load('pretrained_models/s3d_howto100m.pth'))
|
||||
model = model.cuda()
|
||||
|
||||
elif args.type == '3d':
|
||||
print('Loading 3D-ResneXt-101 ...')
|
||||
from videocnn.models import resnext
|
||||
model = resnext.resnet101(
|
||||
num_classes=400,
|
||||
shortcut_type='B',
|
||||
cardinality=32,
|
||||
sample_size=112,
|
||||
sample_duration=16,
|
||||
last_fc=False)
|
||||
model = model.cuda()
|
||||
model_data = th.load(args.resnext101_model_path)
|
||||
model.load_state_dict(model_data)
|
||||
elif args.type == 'vae':
|
||||
from openaivae import OpenAIParallelDiscreteVAE
|
||||
model = OpenAIParallelDiscreteVAE()
|
||||
model = model.cuda()
|
||||
else:
|
||||
raise ValueError("model not supported yet.")
|
||||
|
||||
model.eval()
|
||||
print('loaded')
|
||||
return model
|
||||
@@ -0,0 +1,89 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import os
|
||||
import urllib.parse
|
||||
import json
|
||||
import pandas as pd
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
# TODO: extending to other datasets.
|
||||
supported_formats = {}
|
||||
|
||||
|
||||
class PathBuilder(object):
|
||||
@classmethod
|
||||
def build(cls, video_dirs, feature_dir, ext, shards=0, split=None):
|
||||
meta_fn = os.path.join(feature_dir, "meta_plan.json")
|
||||
os.makedirs(feature_dir, exist_ok=True)
|
||||
if os.path.isfile(meta_fn):
|
||||
with open(meta_fn) as fr:
|
||||
meta = json.load(fr)
|
||||
return meta
|
||||
print("searching videos...")
|
||||
|
||||
video_id_to_path = {}
|
||||
for video_dir in video_dirs.split(","):
|
||||
# TODO: add supports of recursive listdir.
|
||||
if video_dir in supported_formats:
|
||||
supported_formats[video_dir].load(video_dir, video_id_to_path)
|
||||
else:
|
||||
for idx, fn in enumerate(tqdm(os.listdir(video_dir))):
|
||||
video_fn = os.path.join(video_dir, fn)
|
||||
if os.path.isfile(video_fn):
|
||||
video_id = os.path.splitext(fn)[0]
|
||||
video_id_to_path[video_id] = video_fn
|
||||
elif os.path.isdir(video_fn):
|
||||
# shards of folders.
|
||||
shard_dir = video_fn
|
||||
for idx, fn in enumerate(os.listdir(shard_dir)):
|
||||
video_fn = os.path.join(shard_dir, fn)
|
||||
if os.path.isfile(video_fn):
|
||||
video_id = os.path.splitext(fn)[0]
|
||||
video_id_to_path[video_id] = video_fn
|
||||
|
||||
video_path, feature_path = [], []
|
||||
valid_ext = set()
|
||||
for idx, video_id in enumerate(video_id_to_path):
|
||||
video_path.append(video_id_to_path[video_id])
|
||||
if ext is None:
|
||||
# use original file ext for format compatibility.
|
||||
video_id_to_path[video_id]
|
||||
path = urllib.parse.urlparse(video_id_to_path[video_id]).path
|
||||
ext = os.path.splitext(path)[1]
|
||||
if ext not in valid_ext:
|
||||
valid_ext.add(ext)
|
||||
print("adding", ext)
|
||||
if shards:
|
||||
shard_id = str(idx % shards)
|
||||
feature_fn = os.path.join(
|
||||
feature_dir, shard_id, video_id + ext)
|
||||
else:
|
||||
feature_fn = os.path.join(
|
||||
feature_dir, video_id + ext)
|
||||
feature_path.append(feature_fn)
|
||||
|
||||
print("targeting", len(feature_path), "videos")
|
||||
meta = {
|
||||
"video_path": video_path, "feature_path": feature_path}
|
||||
with open(meta_fn, "w") as fw:
|
||||
json.dump(meta, fw)
|
||||
|
||||
if split is not None:
|
||||
splits = split.split("/")
|
||||
assert len(splits) == 2
|
||||
cur, total = int(splits[0]), int(splits[1])
|
||||
assert cur < total
|
||||
import math
|
||||
chunk = math.ceil(len(meta["video_path"]) / total)
|
||||
start = cur * chunk
|
||||
end = (cur + 1) * chunk
|
||||
meta = {
|
||||
"video_path": meta["video_path"][start:end],
|
||||
"feature_path": meta["feature_path"][start:end]
|
||||
}
|
||||
|
||||
return meta
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright Howto100m authors.
|
||||
# Copyright (c) Facebook, Inc. All Rights Reserved
|
||||
|
||||
import torch as th
|
||||
|
||||
class Normalize(object):
|
||||
|
||||
def __init__(self, mean, std):
|
||||
self.mean = th.FloatTensor(mean).view(1, 3, 1, 1)
|
||||
self.std = th.FloatTensor(std).view(1, 3, 1, 1)
|
||||
|
||||
def __call__(self, tensor):
|
||||
tensor = (tensor - self.mean) / (self.std + 1e-8)
|
||||
return tensor
|
||||
|
||||
class Preprocessing(object):
|
||||
|
||||
def __init__(self, type):
|
||||
self.type = type
|
||||
if type == '2d':
|
||||
self.norm = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
|
||||
elif type == '3d':
|
||||
self.norm = Normalize(mean=[110.6, 103.2, 96.3], std=[1.0, 1.0, 1.0])
|
||||
elif type == 'vmz':
|
||||
self.norm = Normalize(mean=[110.201, 100.64, 95.997], std=[58.1489, 56.4701, 55.3324])
|
||||
|
||||
def _zero_pad(self, tensor, size):
|
||||
n = size - len(tensor) % size
|
||||
if n == size:
|
||||
return tensor
|
||||
else:
|
||||
z = th.zeros(n, tensor.shape[1], tensor.shape[2], tensor.shape[3])
|
||||
return th.cat((tensor, z), 0)
|
||||
|
||||
def __call__(self, tensor):
|
||||
if self.type == '2d':
|
||||
tensor = tensor / 255.0
|
||||
tensor = self.norm(tensor)
|
||||
elif self.type == 'vmz':
|
||||
#tensor = self._zero_pad(tensor, 8)
|
||||
tensor = self._zero_pad(tensor, 10)
|
||||
tensor = self.norm(tensor)
|
||||
#tensor = tensor.view(-1, 8, 3, 112, 112)
|
||||
tensor = tensor.view(-1, 10, 3, 112, 112)
|
||||
tensor = tensor.transpose(1, 2)
|
||||
elif self.type == '3d':
|
||||
tensor = self._zero_pad(tensor, 16)
|
||||
tensor = self.norm(tensor)
|
||||
tensor = tensor.view(-1, 16, 3, 112, 112)
|
||||
tensor = tensor.transpose(1, 2)
|
||||
elif self.type == 's3d':
|
||||
tensor = tensor / 255.0
|
||||
tensor = self._zero_pad(tensor, 30)
|
||||
tensor = tensor.view(-1, 30, 3, 224, 224) # N x 30 x 3 x H x W
|
||||
tensor = tensor.transpose(1, 2) # N x 3 x 30 x H x W
|
||||
# for vae do nothing
|
||||
return tensor
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) Facebook, Inc. All Rights Reserved
|
||||
|
||||
import numpy as np
|
||||
|
||||
from torch.utils.data.sampler import Sampler
|
||||
|
||||
|
||||
class RandomSequenceSampler(Sampler):
|
||||
|
||||
def __init__(self, n_sample, seq_len):
|
||||
self.n_sample = n_sample
|
||||
self.seq_len = seq_len
|
||||
|
||||
def _pad_ind(self, ind):
|
||||
zeros = np.zeros(self.seq_len - self.n_sample % self.seq_len)
|
||||
ind = np.concatenate((ind, zeros))
|
||||
return ind
|
||||
|
||||
def __iter__(self):
|
||||
idx = np.arange(self.n_sample)
|
||||
if self.n_sample % self.seq_len != 0:
|
||||
idx = self._pad_ind(idx)
|
||||
idx = np.reshape(idx, (-1, self.seq_len))
|
||||
np.random.shuffle(idx)
|
||||
idx = np.reshape(idx, (-1))
|
||||
return iter(idx.astype(int))
|
||||
|
||||
def __len__(self):
|
||||
return self.n_sample + (self.seq_len - self.n_sample % self.seq_len)
|
||||
@@ -0,0 +1,64 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import numpy as np
|
||||
import os
|
||||
import pickle
|
||||
|
||||
from mmpt.utils import ShardedTensor
|
||||
|
||||
|
||||
class Shard(object):
|
||||
def __init__(
|
||||
self,
|
||||
vfeat_dir,
|
||||
tfeat_dir,
|
||||
target_dir,
|
||||
file_paths,
|
||||
shard_size=4096
|
||||
):
|
||||
self.vfeat_dir = vfeat_dir
|
||||
self.tfeat_dir = tfeat_dir
|
||||
self.target_dir = target_dir
|
||||
self.video_ids = {}
|
||||
for split, file_path in zip(["train", "val"], file_paths):
|
||||
with open(file_path) as fr:
|
||||
self.video_ids[split] = [
|
||||
line.strip() for line in fr.readlines()]
|
||||
self.shard_size = shard_size
|
||||
|
||||
def __call__(self, split="train"):
|
||||
for split in ["train", "val"]:
|
||||
meta = {}
|
||||
for shard_idx, shard_offset in enumerate(
|
||||
range(0, len(self.video_ids[split]), self.shard_size)
|
||||
):
|
||||
print(shard_idx)
|
||||
meta_shard = []
|
||||
video_shard = []
|
||||
for video_id in self.video_ids[split][shard_offset:shard_offset+self.shard_size]:
|
||||
meta_shard.append(video_id)
|
||||
npy_file = os.path.join(self.vfeat_dir, video_id + ".npy")
|
||||
video_shard.append(np.load(npy_file))
|
||||
|
||||
meta[shard_idx] = meta_shard
|
||||
video_shard = ShardedTensor.from_list(video_shard)
|
||||
target_path = os.path.join(
|
||||
self.target_dir, split + "_" + str(shard_idx))
|
||||
video_shard.save(target_path)
|
||||
|
||||
target_path = os.path.join(self.target_dir, split + "_meta")
|
||||
with open(target_path + ".pkl", "wb") as fw:
|
||||
pickle.dump(meta, fw, pickle.HIGHEST_PROTOCOL)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
shard = Shard(
|
||||
"data/feat/feat_how2_s3d",
|
||||
"data/how2/raw_caption_dedup.bert-base-uncased",
|
||||
"data/feat/feat_how2_s3d_shard_small",
|
||||
["data/how2/how2_s3d_train.lst", "data/how2/how2_s3d_val.lst"]
|
||||
)
|
||||
|
||||
shard()
|
||||
@@ -0,0 +1,242 @@
|
||||
# Copyright Howto100M authors.
|
||||
# Copyright (c) Facebook, Inc. All Rights Reserved
|
||||
|
||||
import torch as th
|
||||
import pandas as pd
|
||||
import os
|
||||
import numpy as np
|
||||
import ffmpeg
|
||||
import random
|
||||
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class VideoLoader(Dataset):
|
||||
"""modified from how2's video_feature_extractor."""
|
||||
def __init__(
|
||||
self,
|
||||
csv=None,
|
||||
video_dict=None,
|
||||
framerate=1,
|
||||
size=112,
|
||||
centercrop=False,
|
||||
hflip=False,
|
||||
**kwargs
|
||||
):
|
||||
if csv is None and video_dict is None:
|
||||
raise ValueError("csv and video_dict cannot be both None.")
|
||||
if csv is not None:
|
||||
self.csv = pd.read_csv(csv)
|
||||
if video_dict is not None:
|
||||
self.csv = pd.DataFrame.from_dict(video_dict)
|
||||
|
||||
self.centercrop = centercrop
|
||||
self.size = size
|
||||
self.framerate = framerate
|
||||
self.hflip = hflip
|
||||
|
||||
def __len__(self):
|
||||
return len(self.csv)
|
||||
|
||||
def _get_video_dim(self, video_path):
|
||||
probe = ffmpeg.probe(video_path)
|
||||
video_stream = next((stream for stream in probe['streams']
|
||||
if stream['codec_type'] == 'video'), None)
|
||||
width = int(video_stream['width'])
|
||||
height = int(video_stream['height'])
|
||||
return height, width
|
||||
|
||||
def _get_video_info(self, video_path):
|
||||
probe = ffmpeg.probe(video_path)
|
||||
video_stream = next((stream for stream in probe['streams']
|
||||
if stream['codec_type'] == 'video'), None)
|
||||
return video_stream
|
||||
|
||||
def _get_output_dim(self, h, w):
|
||||
if isinstance(self.size, tuple) and len(self.size) == 2:
|
||||
return self.size
|
||||
elif h >= w:
|
||||
return int(h * self.size / w), self.size
|
||||
else:
|
||||
return self.size, int(w * self.size / h)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
video_path = self.csv['video_path'].values[idx]
|
||||
output_file = self.csv['feature_path'].values[idx]
|
||||
return self._decode(output_file, video_path)
|
||||
|
||||
def _decode(self, output_file, video_path):
|
||||
if not(os.path.isfile(output_file)) and os.path.isfile(video_path):
|
||||
try:
|
||||
h, w = self._get_video_dim(video_path)
|
||||
except Exception:
|
||||
print('ffprobe failed at: {}'.format(video_path))
|
||||
return {'video': th.zeros(1), 'input': video_path,
|
||||
'output': output_file}
|
||||
try:
|
||||
os.makedirs(os.path.dirname(output_file), exist_ok=True)
|
||||
height, width = self._get_output_dim(h, w)
|
||||
|
||||
cmd = (
|
||||
ffmpeg
|
||||
.input(video_path)
|
||||
.filter('fps', fps=self.framerate)
|
||||
.filter('scale', width, height)
|
||||
)
|
||||
if self.hflip:
|
||||
cmd = cmd.filter('hflip')
|
||||
|
||||
if self.centercrop:
|
||||
x = int((width - self.size) / 2.0)
|
||||
y = int((height - self.size) / 2.0)
|
||||
cmd = cmd.crop(x, y, self.size, self.size)
|
||||
video = self._run(cmd, output_file)
|
||||
except Exception:
|
||||
video = th.zeros(1)
|
||||
else:
|
||||
video = th.zeros(1)
|
||||
|
||||
return {'video': video, 'input': video_path, 'output': output_file}
|
||||
|
||||
def _run(self, cmd, output_file):
|
||||
out, _ = (
|
||||
cmd.output('pipe:', format='rawvideo', pix_fmt='rgb24')
|
||||
.run(capture_stdout=True, quiet=True)
|
||||
)
|
||||
if self.centercrop and isinstance(self.size, int):
|
||||
height, width = self.size, self.size
|
||||
video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
|
||||
video = th.from_numpy(video.astype('float32'))
|
||||
return video.permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
class VideoVerifier(VideoLoader):
|
||||
def __getitem__(self, idx):
|
||||
video_path = self.csv['video_path'].values[idx]
|
||||
try:
|
||||
return self._get_video_info(video_path)
|
||||
except Exception:
|
||||
# print('ffprobe failed at: {}'.format(video_path))
|
||||
return None
|
||||
|
||||
|
||||
class VideoCompressor(VideoLoader):
|
||||
def __init__(
|
||||
self,
|
||||
csv=None,
|
||||
video_dict=None,
|
||||
framerate=1,
|
||||
size=112,
|
||||
centercrop=False,
|
||||
hflip=False,
|
||||
crf=32,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(
|
||||
csv,
|
||||
video_dict,
|
||||
framerate,
|
||||
size,
|
||||
centercrop,
|
||||
hflip
|
||||
)
|
||||
self.crf = crf
|
||||
|
||||
def _run(self, cmd, output_file):
|
||||
out, _ = (
|
||||
cmd.output(filename=output_file, crf=self.crf)
|
||||
.run(quiet=True)
|
||||
)
|
||||
video = None
|
||||
return video
|
||||
|
||||
|
||||
class VideoDownloader(VideoCompressor):
|
||||
"""download"""
|
||||
def __getitem__(self, idx):
|
||||
video_path = self.csv['video_path'].values[idx]
|
||||
output_file = self.csv['feature_path'].values[idx]
|
||||
if not(os.path.isfile(output_file)):
|
||||
os.makedirs(os.path.dirname(output_file), exist_ok=True)
|
||||
cmd = "wget -O" + output_file + " " + video_path
|
||||
# import subprocess
|
||||
# subprocess.check_output(
|
||||
# cmd,
|
||||
# stderr=subprocess.STDOUT, shell=True)
|
||||
os.system(cmd)
|
||||
return {'video': None, 'input': video_path, 'output': output_file}
|
||||
|
||||
|
||||
class AvKeyframeVideoCompressor(VideoLoader):
|
||||
"""extract keyframes from a video and save it as jpg.
|
||||
TODO: consider to merge with `CodecProcessor`.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
csv=None,
|
||||
video_dict=None,
|
||||
framerate=1,
|
||||
size=112,
|
||||
centercrop=False,
|
||||
max_num_frames=5,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(csv, video_dict, framerate, size, centercrop)
|
||||
self.max_num_frames = max_num_frames
|
||||
|
||||
def _get_video_dim(self, video_fn):
|
||||
"""decord cannot probe the size of a video, we use pyav instead."""
|
||||
import av
|
||||
with av.open(video_fn) as container:
|
||||
height = container.streams.video[0].codec_context.height
|
||||
width = container.streams.video[0].codec_context.width
|
||||
return height, width
|
||||
|
||||
def _get_output_dim(self, height, width):
|
||||
"""
|
||||
keep the shorter side be `self.size`, strech the other.
|
||||
"""
|
||||
if height >= width:
|
||||
return int(height * self.size / width), self.size
|
||||
else:
|
||||
return self.size, int(width * self.size / height)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
import av
|
||||
video_path = self.csv['video_path'].values[idx]
|
||||
output_file = self.csv['feature_path'].values[idx]
|
||||
if not(os.path.isdir(output_file)) and os.path.isfile(video_path):
|
||||
try:
|
||||
h, w = self._get_video_dim(video_path)
|
||||
except Exception:
|
||||
print('probe failed at: {}'.format(video_path))
|
||||
return {'video': th.zeros(1), 'input': video_path,
|
||||
'output': output_file}
|
||||
|
||||
try:
|
||||
height, width = self._get_output_dim(h, w)
|
||||
|
||||
# new for av.
|
||||
with av.open(video_path) as container:
|
||||
container.streams.video[0].thread_type = "AUTO"
|
||||
container.streams.video[0].codec_context.height = height
|
||||
container.streams.video[0].codec_context.width = width
|
||||
if self.framerate == 0: # keyframe.
|
||||
container.streams.video[0].codec_context.skip_frame = 'NONKEY'
|
||||
frames = []
|
||||
for frame in container.decode(video=0):
|
||||
frames.append(frame)
|
||||
frames = random.sample(frames, self.max_num_frames)
|
||||
|
||||
os.makedirs(output_file, exist_ok=True)
|
||||
for frame in frames:
|
||||
frame.to_image().save(
|
||||
os.path.join(
|
||||
output_file,
|
||||
"%04d.jpg" % frame.index))
|
||||
except Exception:
|
||||
print('extract failed at: {}'.format(video_path))
|
||||
return {'video': th.zeros(1), 'input': video_path,
|
||||
'output': output_file}
|
||||
video = th.zeros(1)
|
||||
return {'video': video, 'input': video_path, 'output': output_file}
|
||||
Reference in New Issue
Block a user