Update image process (#5)

* [docs] update tool docs

* update aes

README.md

@@ -38,7 +38,7 @@ With Open-Sora, we aim to inspire innovation, creativity, and inclusivity in the
 ## 🎥 Latest Demo
 
 | **2s 512×512**                                                                                                                                                                 | **2s 512×512**                                                                                                                                                              | **2s 512×512**                                                                                                                                    |
-|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------|
+| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------- |
 | [<img src="assets/readme/sample_0.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/de1963d3-b43b-4e68-a670-bb821ebb6f80)                                 | [<img src="assets/readme/sample_1.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/13f8338f-3d42-4b71-8142-d234fbd746cc)                              | [<img src="assets/readme/sample_2.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/fa6a65a6-e32a-4d64-9a9e-eabb0ebb8c16)    |
 | A serene night scene in a forested area. [...] The video is a time-lapse, capturing the transition from day to night, with the lake and forest serving as a constant backdrop. | A soaring drone footage captures the majestic beauty of a coastal cliff, [...] The water gently laps at the rock base and the greenery that clings to the top of the cliff. | The majestic beauty of a waterfall cascading down a cliff into a serene lake. [...] The camera angle provides a bird's eye view of the waterfall. |
 | [<img src="assets/readme/sample_3.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/64232f84-1b36-4750-a6c0-3e610fa9aa94)                                 | [<img src="assets/readme/sample_4.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/983a1965-a374-41a7-a76b-c07941a6c1e9)                              | [<img src="assets/readme/sample_5.gif" width="">](https://github.com/hpcaitech/Open-Sora/assets/99191637/ec10c879-9767-4c31-865f-2e8d6cf11e65)    |
@@ -144,7 +144,7 @@ the config files.
 ## Model Weights
 
 | Resolution | Data   | #iterations | Batch Size | GPU days (H800) | URL                                                                                           |
-|------------|--------|-------------|------------|-----------------|-----------------------------------------------------------------------------------------------|
+| ---------- | ------ | ----------- | ---------- | --------------- | --------------------------------------------------------------------------------------------- |
 | 16×512×512 | 20K HQ | 20k         | 2×64       | 35              | [:link:](https://huggingface.co/hpcai-tech/Open-Sora/blob/main/OpenSora-v1-HQ-16x512x512.pth) |
 | 16×256×256 | 20K HQ | 24k         | 8×64       | 45              | [:link:](https://huggingface.co/hpcai-tech/Open-Sora/blob/main/OpenSora-v1-HQ-16x256x256.pth) |
 | 16×256×256 | 366K   | 80k         | 8×64       | 117             | [:link:](https://huggingface.co/hpcai-tech/Open-Sora/blob/main/OpenSora-v1-16x256x256.pth)    |
@@ -254,6 +254,9 @@ following [all-contributors](https://github.com/all-contributors/all-contributor
 
 If you wish to contribute to this project, you can refer to the [Contribution Guideline](./CONTRIBUTING.md).
 
+[Zangwei Zheng](https://github.com/zhengzangw) and [Xiangyu Peng](https://github.com/xyupeng) equally contributed to
+this work during their internship at [HPC-AI Tech](https://hpc-ai.com/).
+
 ## Acknowledgement
 
 * [ColossalAI](https://github.com/hpcaitech/ColossalAI): A powerful large model parallel acceleration and optimization
@@ -271,7 +274,7 @@ If you wish to contribute to this project, you can refer to the [Contribution Gu
 
 We are grateful for their exceptional work and generous contribution to open source.
 
-## Citation
+<!-- ## Citation
 
 ```bibtex
 @software{opensora,
@@ -281,10 +284,7 @@ We are grateful for their exceptional work and generous contribution to open sou
   year = {2024},
   url = {https://github.com/hpcaitech/Open-Sora}
 }
-```
+``` -->
-
-[Zangwei Zheng](https://github.com/zhengzangw) and [Xiangyu Peng](https://github.com/xyupeng) equally contributed to
-this work during their internship at [HPC-AI Tech](https://hpc-ai.com/).
 
 ## Star History
 

docs/datasets.md

@@ -1,28 +1,15 @@
 # Datasets
 
-## Datasets used for now
+## HD-VG-130M
 
-### HD-VG-130M
+[HD-VG-130M](https://github.com/daooshee/HD-VG-130M?tab=readme-ov-file) comprises 130M text-video pairs. The caption is generated by BLIP-2. We find the cut and the text quality are relatively poor. It contains 20 splits. For OpenSora 1.0, we use the first split (~350K). We plan to use the whole dataset and re-process it.
 
-[HD-VG-130M](https://github.com/daooshee/HD-VG-130M?tab=readme-ov-file) comprises 130M text-video pairs. The caption is generated by BLIP-2. We find the cut and the text quality are relatively poor. It contains 20 splits. For OpenSora 1.0, we use the first split. We plan to use the whole dataset and re-process it.
+You can download the dataset and prepare it for training according to [the dataset repository's instructions](https://github.com/daooshee/HD-VG-130M). There is a README.md file in the Google Drive link that provides instructions on how to download and cut the videos. For this version, we directly use the dataset provided by the authors.
 
-### Inter4k
+## Inter4k
 
 [Inter4k](https://github.com/alexandrosstergiou/Inter4K) is a dataset containing 1k video clips with 4K resolution. The dataset is proposed for super-resolution tasks. We use the dataset for HQ training. The videos are processed as mentioned [here](/README.md#data-processing).
 
-### Pexels.com
+## Pexels.com
 
 [Pexels.com](https://www.pexels.com/) is a website that provides free stock photos and videos. We collect 19K video clips from this website for HQ training. The videos are processed as mentioned [here](/README.md#data-processing).
-
-## Datasets watching list
-
-We are also watching the following datasets and considering using them in the future, which depends on our disk space and the quality of the dataset.
-
-| Name              | Size         | Description                   |
-| ----------------- | ------------ | ----------------------------- |
-| Panda-70M         | 70M videos   | High quality video-text pairs |
-| WebVid-10M        | 10M videos   | Low quality                   |
-| InternVid-10M-FLT | 10M videos   |                               |
-| EGO4D             | 3670 hours   |                               |
-| OpenDV-YouTube    | 1700 hours   |                               |
-| VidProM           | 6.69M videos |                               |

tools/aesthetic/README.md

@@ -1,6 +1,26 @@
 # Aesthetic Scoring
 
+To evaluate the aesthetic quality of videos, we use a pretrained model from [CLIP+MLP Aesthetic Score Predictor](https://github.com/christophschuhmann/improved-aesthetic-predictor). This model is trained on 176K SAC (Simulacra Aesthetic Captions) pairs, 15K LAION-Logos (Logos) pairs, and 250K AVA (The Aesthetic Visual Analysis) image-text pairs.
+
+The score is between 1 and 10, where 5.5 can be considered as the threshold for fair aesthetics, and 6.5 for good aesthetics. Good text-to-image models can achieve a score of 7.0 or higher.
+
+For videos, we extract the first, last, and the middle frames for evaluation. The script also supports images. Our script enables 1k videos/s with one GPU. It also supports multiple GPUs to further accelerate the process.
+
+## Requirement
+
 ```bash
+# install clip
 pip install git+https://github.com/openai/CLIP.git
+
+# get pretrained model
 wget https://github.com/christophschuhmann/improved-aesthetic-predictor/raw/main/sac+logos+ava1-l14-linearMSE.pth -O pretrained_models/aesthetic.pth
 ```
+
+## Usage
+
+With `DATA.csv` containing the paths to the videos, run the following command:
+
+```bash
+# output: DATA_aes.csv
+python -m tools.aesthetic.inference DATA.csv
+```

tools/aesthetic/inference.py

@@ -1,5 +1,6 @@
 # adapted from https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/simple_inference.py
 import argparse
+import os
 
 import av
 import clip
@@ -9,10 +10,19 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
+from PIL import Image
 from tqdm import tqdm
 
+IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp")
+VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv")
 
-def extract_frames(video_path, points=(0.0, 0.5, 0.9)):
+
+def is_video(filename):
+    ext = os.path.splitext(filename)[-1].lower()
+    return ext in VID_EXTENSIONS
+
+
+def extract_frames(video_path, points=(0.1, 0.5, 0.9)):
     container = av.open(video_path)
     total_frames = container.streams.video[0].frames
     frames = []
@@ -25,6 +35,10 @@ def extract_frames(video_path, points=(0.0, 0.5, 0.9)):
     return frames
 
 
+def get_image(image_path):
+    return Image.open(image_path).convert("RGB")
+
+
 class VideoTextDataset(torch.utils.data.Dataset):
     def __init__(self, csv_path, transform=None, points=(0.1, 0.5, 0.9)):
         self.csv_path = csv_path
@@ -34,7 +48,11 @@ class VideoTextDataset(torch.utils.data.Dataset):
 
     def getitem(self, index):
         sample = self.data.iloc[index]
-        images = extract_frames(sample["path"], points=self.points)
+        path = sample["path"]
+        if not is_video(path):
+            images = [get_image(path)]
+        else:
+            images = extract_frames(sample["path"], points=self.points)
         images = [self.transform(img) for img in images]
         images = torch.stack(images)
 
@@ -117,6 +135,7 @@ def main(args):
         dataset.data.loc[index : index + len(scores_np) - 1, "aesthetic"] = scores_np
         index += len(images)
     dataset.data.to_csv(output_file, index=False)
+    print(f"Saved aesthetic scores to {output_file}.")
 
 
 if __name__ == "__main__":
@@ -126,4 +145,5 @@ if __name__ == "__main__":
     parser.add_argument("--num_workers", type=int, default=64, help="Number of workers")
     parser.add_argument("--prefetch_factor", type=int, default=8, help="Prefetch factor")
     args = parser.parse_args()
+
     main(args)

tools/datasets/README.md

@@ -1,31 +1,10 @@
-# Dataset Download and Management
+# Dataset Management
 
-## Dataset Download
+After preparing the raw dataset according to the [instructions](/docs/datasets.md), you can use the following commands to manage the dataset.
-
-### HD-VG-130M
-
-This dataset comprises 130M text-video pairs. You can download the dataset and prepare it for training according to [the dataset repository's instructions](https://github.com/daooshee/HD-VG-130M). There is a README.md file in the Google Drive link that provides instructions on how to download and cut the videos. For this version, we directly use the dataset provided by the authors.
-
-### VidProM
-
-```bash
-python -m tools.datasets.convert_dataset vidprom VIDPROM_FOLDER --info VidProM_semantic_unique.csv
-```
-
-### Demo Dataset
-
-You can use ImageNet and UCF101 for a quick demo. After downloading the datasets, you can use the following command to prepare the csv file for the dataset:
-
-```bash
-# ImageNet
-python -m tools.datasets.convert_dataset imagenet IMAGENET_FOLDER --split train
-# UCF101
-python -m tools.datasets.convert_dataset ucf101 UCF101_FOLDER --split videos
-```
 
 ## Dataset Format
 
-The 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). Aspect ratio is width divided by height.
+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).
 
 ```csv
 path, text, num_frames, fps, width, height, aspect_ratio, aesthetic_score, clip_score
@@ -33,50 +12,151 @@ path, text, num_frames, fps, width, height, aspect_ratio, aesthetic_score, clip_
 /absolute/path/to/video2.mp4, caption2, num_of_frames
 ```
 
-We use pandas to manage the CSV files. You can use the following code to read and write the CSV files:
+We use pandas to manage the CSV files. The following code is for reading and writing the CSV files:
 
 ```python
 df = pd.read_csv(input_path)
 df = df.to_csv(output_path, index=False)
 ```
 
-## Manage datasets
+The columns are defined as follows:
 
-We provide `csvutils.py` to manage the CSV files. You can use the following commands to process the CSV files:
+- `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](/tools/aesthetic/README.md). Optional.
+- `clip_score`: the clip score by [clip scorer](/tools/clip/README.md). 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:
 
 ```bash
-# csvutil takes multiple CSV files as input and merge them into one CSV file
+python -m tools.datasets.convert_dataset DATASET-TYPE DATA_FOLDER
-python -m tools.datasets.csvutil DATA1.csv DATA2.csv
+# 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
+```
 
-# filter frames between 128 and 256, with captions
+## Manage datasets
-python -m tools.datasets.csvutil DATA.csv --fmin 128 --fmax 256 --remove-empty-caption
-# compute the number of frames for each video
-python -m tools.datasets.csvutil DATA.csv --video-info
-# remove caption prefix
-python -m tools.datasets.csvutil DATA.csv --remove-caption-prefix
-# generate DATA_root.csv with absolute path
-python -m tools.datasets.csvutil DATA.csv --abspath /absolute/path/to/dataset
 
+You can easily get basic information about the dataset by using the following commands:
+
+```bash
 # 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](https://github.com/nalepae/pandarallel):
 
 ```bash
 pip install pandarallel
 ```
 
-To filter text language, you need to install [lingua](https://github.com/pemistahl/lingua-py):
-
-```bash
-pip install lingua-language-detector
-```
-
 To get video information, you need to install [opencv-python](https://github.com/opencv/opencv-python):
 
 ```bash
 pip install opencv-python
 ```
+
+To filter a specific language, you need to install [lingua](https://github.com/pemistahl/lingua-py):
+
+```bash
+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.
+
+```bash
+# 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:
+
+```bash
+# 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:
+
+```bash
+# 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:
+
+```bash
+# 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
+```

tools/datasets/convert_dataset.py

@@ -6,7 +6,7 @@ from torchvision.datasets import ImageNet
 
 
 IMG_EXTENSIONS = (".jpg", ".jpeg", ".png", ".ppm", ".bmp", ".pgm", ".tif", ".tiff", ".webp")
-VID_EXTENSIONS = ("mp4", "avi", "mov", "mkv")
+VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv")
 
 
 def get_filelist(file_path, exts=None):

tools/datasets/csvutil.py

@@ -131,7 +131,7 @@ def get_output_path(args, input_name):
         name += "_relpath"
     # caption filtering
     if args.remove_empty_caption:
-        name += "_rec"
+        name += "_noempty"
     if args.lang is not None:
         name += f"_{args.lang}"
     if args.remove_url: