Open-Sora/tools/datasets

Dataset Management

After preparing the raw dataset according to the instructions, you can use the following commands to manage the dataset.

Dataset Format

All dataset should be provided in a CSV file, which is used both for training and data preprocessing.. The CSV file should only contain the following columns (can be optional).

path, text, num_frames, fps, width, height, aspect_ratio, aesthetic_score, clip_score
/absolute/path/to/image1.jpg, caption1, num_of_frames
/absolute/path/to/video2.mp4, caption2, num_of_frames

We use pandas to manage the CSV files. The following code is for reading and writing the CSV files:

df = pd.read_csv(input_path)
df = df.to_csv(output_path, index=False)

The columns are defined as follows:

• path: the relative/absolute path or url to the image or video file. The only required column.
• text: the caption or description of the image or video. Necessary for training.
• num_frames: the number of frames in the video. Necessary for training.
• fps: the frame rate of the video. Optional.
• width: the width of the video frame. Necessary for STDiT2.
• height: the height of the video frame. Necessary for STDiT2.
• aspect_ratio: the aspect ratio of the video frame (height divided by width). Optional.
• aesthetic_score: the aesthetic score by asethetic scorer. Optional.
• clip_score: the clip score by clip scorer. Optional.

Dataset to CSV

As a start point, convert_dataset.py is used to convert the dataset to a CSV file. You can use the following commands to convert the dataset to a CSV file:

python -m tools.datasets.convert_dataset DATASET-TYPE DATA_FOLDER
# general video folder
python -m tools.datasets.convert_dataset video VIDEO_FOLDER
# general image folder
python -m tools.datasets.convert_dataset image IMAGE_FOLDER
# imagenet
python -m tools.datasets.convert_dataset imagenet IMAGENET_FOLDER --split train
# ucf101
python -m tools.datasets.convert_dataset ucf101 UCF101_FOLDER --split videos
# vidprom
python -m tools.datasets.convert_dataset vidprom VIDPROM_FOLDER --info VidProM_semantic_unique.csv

Manage datasets

You can easily get basic information about the dataset by using the following commands:

# examine the first 10 rows of the CSV file
head -n 10 DATA1.csv
# count the number of data in the CSV file (approximately)
wc -l DATA1.csv

Additionally, Ww provide csvutils.py to manage the CSV files.

Requirement

To accelerate processing speed, you can install pandarallel:

pip install pandarallel

To get video information, you need to install opencv-python:

pip install opencv-python

To filter a specific language, you need to install lingua:

pip install lingua-language-detector

Usage

You can use the following commands to process the CSV files. The output csv file will be saved in the same directory as the input csv file, with different suffixes indicating the processing method.

# csvutil takes multiple CSV files as input and merge them into one CSV file
# output: DATA1+DATA2.csv
python -m tools.datasets.csvutil DATA1.csv DATA2.csv
# shard CSV files into multiple CSV files
# output: DATA1_0.csv, DATA1_1.csv, ...
python -m tools.datasets.csvutil DATA1.csv --shard 10
# filter frames between 128 and 256, with captions
# output: DATA1_fmin_128_fmax_256.csv
python -m tools.datasets.csvutil DATA.csv --fmin 128 --fmax 256
# Disable parallel processing
python -m tools.datasets.csvutil DATA.csv --fmin 128 --fmax 256 --disable-parallel

Here are more examples:

# modify the path to absolute path by root given
# output: DATA_abspath.csv
python -m tools.datasets.csvutil DATA.csv --abspath /absolute/path/to/dataset
# modify the path to relative path by root given
# output: DATA_relpath.csv
python -m tools.datasets.csvutil DATA.csv --relpath /relative/path/to/dataset

# remove the rows with empty captions
# output: DATA_noempty.csv
python -m tools.datasets.csvutil DATA.csv --remove-empty-caption
# remove the rows with urls
# output: DATA_nourl.csv
python -m tools.datasets.csvutil DATA.csv --remove-url
# unescape the caption
# output: DATA_unescape.csv
python -m tools.datasets.csvutil DATA.csv --unescape
# modify LLaVA caption
# output: DATA_rcp.csv
python -m tools.datasets.csvutil DATA.csv --remove-caption-prefix
# keep only the rows with english captions
# output: DATA_en.csv
python -m tools.datasets.csvutil DATA.csv --lang en

# compute num_frames, height, width, fps, aspect_ratio for videos or images
# output: IMG_DATA+VID_DATA_vinfo.csv
python -m tools.datasets.csvutil IMG_DATA.csv VID_DATA --video-info

You can apply multiple operations at the same time:

# output: DATA_vinfo_noempty_nourl_en.csv
python -m tools.datasets.csvutil DATA.csv --video-info --remove-empty-caption --remove-url --lang en

To examine and filter the quality of the dataset by aesthetic score and clip score, you can use the following commands:

# sort the dataset by aesthetic score
# output: DATA_sort.csv
python -m tools.datasets.csvutil DATA.csv --sort-descending aesthetic_score
# View examples of high aesthetic score
head -n 10 DATA_sort.csv
# View examples of low aesthetic score
tail -n 10 DATA_sort.csv

# sort the dataset by clip score
# output: DATA_sort.csv
python -m tools.datasets.csvutil DATA.csv --sort-descending clip_score

# filter the dataset by aesthetic score
# output: DATA_aesmin_0.5.csv
python -m tools.datasets.csvutil DATA.csv --aesmin 0.5
# filter the dataset by clip score
# output: DATA_matchmin_0.5.csv
python -m tools.datasets.csvutil DATA.csv --matchmin 0.5