wehub-resource-sync
90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
# 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
|