Spaces:
Runtime error
Runtime error
File size: 2,457 Bytes
5d32408 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import argparse
from typing import List
import pandas as pd
from mmengine.config import Config
from videogen_hub.pipelines.opensora.opensora.datasets.bucket import Bucket
def split_by_bucket(
bucket: Bucket,
input_files: List[str],
output_path: str,
limit: int,
frame_interval: int,
):
print(f"Split {len(input_files)} files into {len(bucket)} buckets")
total_limit = len(bucket) * limit
bucket_cnt = {}
# get all bucket id
for hw_id, d in bucket.ar_criteria.items():
for t_id, v in d.items():
for ar_id in v.keys():
bucket_id = (hw_id, t_id, ar_id)
bucket_cnt[bucket_id] = 0
output_df = None
# split files
for path in input_files:
df = pd.read_csv(path)
if output_df is None:
output_df = pd.DataFrame(columns=df.columns)
for i in range(len(df)):
row = df.iloc[i]
t, h, w = row["num_frames"], row["height"], row["width"]
bucket_id = bucket.get_bucket_id(t, h, w, frame_interval)
if bucket_id is None:
continue
if bucket_cnt[bucket_id] < limit:
bucket_cnt[bucket_id] += 1
output_df = pd.concat([output_df, pd.DataFrame([row])], ignore_index=True)
if len(output_df) >= total_limit:
break
if len(output_df) >= total_limit:
break
assert len(output_df) <= total_limit
if len(output_df) == total_limit:
print(f"All buckets are full ({total_limit} samples)")
else:
print(f"Only {len(output_df)} files are used")
output_df.to_csv(output_path, index=False)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("input", type=str, nargs="+")
parser.add_argument("-o", "--output", required=True)
parser.add_argument("-c", "--config", required=True)
parser.add_argument("-l", "--limit", default=200, type=int)
args = parser.parse_args()
assert args.limit > 0
cfg = Config.fromfile(args.config)
bucket_config = cfg.bucket_config
# rewrite bucket_config
for ar, d in bucket_config.items():
for frames, t in d.items():
p, bs = t
if p > 0.0:
p = 1.0
d[frames] = (p, bs)
bucket = Bucket(bucket_config)
split_by_bucket(bucket, args.input, args.output, args.limit, cfg.dataset.frame_interval)
|