Update video caption (#129)
This commit is contained in:
hkz
+10
-26
@@ -1,4 +1,4 @@
|
||||
FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
|
||||
FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
|
||||
ENV DEBIAN_FRONTEND noninteractive
|
||||
|
||||
RUN rm -r /etc/apt/sources.list.d/
|
||||
@@ -6,36 +6,24 @@ RUN rm -r /etc/apt/sources.list.d/
|
||||
RUN apt-get update -y && apt-get install -y \
|
||||
libgl1 libglib2.0-0 google-perftools \
|
||||
sudo wget git git-lfs vim tig pkg-config libcairo2-dev \
|
||||
telnet curl net-tools iputils-ping wget jq \
|
||||
python3-pip python-is-python3 python3.10-venv tzdata lsof && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
aria2 telnet curl net-tools iputils-ping jq \
|
||||
python3-pip python-is-python3 python3.10-venv tzdata lsof zip tmux
|
||||
RUN apt-get update && \
|
||||
apt-get install -y software-properties-common && \
|
||||
add-apt-repository ppa:ubuntuhandbook1/ffmpeg6 && \
|
||||
apt-get update && \
|
||||
apt-get install -y ffmpeg
|
||||
|
||||
RUN pip3 install --upgrade pip -i https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
# add all extensions
|
||||
RUN apt-get update -y && apt-get install -y zip && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
RUN pip install wandb tqdm GitPython==3.1.32 Pillow==9.5.0 setuptools --upgrade -i https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
RUN pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu118
|
||||
RUN pip install xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu118
|
||||
|
||||
# install VILA (video-caption)
|
||||
RUN wget https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/vila-1.0.0-torch2.4.0-py3-none-any.whl
|
||||
RUN mv vila-1.0.0-torch2.4.0-py3-none-any.whl vila-1.0.0-py3-none-any.whl
|
||||
RUN pip install vila-1.0.0-py3-none-any.whl --extra-index-url https://download.pytorch.org/whl/cu118
|
||||
RUN wget https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/package/flash_attn-2.6.3%2Bcu118torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
|
||||
RUN pip install flash_attn-2.6.3+cu118torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
|
||||
|
||||
# install llm-awq (video-caption)
|
||||
RUN git clone https://github.com/mit-han-lab/llm-awq /root/llm-awq
|
||||
WORKDIR /root/llm-awq
|
||||
RUN pip install -e .
|
||||
WORKDIR /root/llm-awq/awq/kernels
|
||||
# https://github.com/mit-han-lab/llm-awq/issues/93#issuecomment-2144434686
|
||||
RUN export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0+PTX" && python3 setup.py install
|
||||
|
||||
# install vllm (video-caption)
|
||||
RUN pip install https://github.com/vllm-project/vllm/releases/download/v0.5.4/vllm-0.5.4+cu118-cp310-cp310-manylinux1_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cu118
|
||||
RUN pip install vllm==0.6.3
|
||||
|
||||
# install requirements (video-caption)
|
||||
WORKDIR /root/
|
||||
@@ -43,7 +31,6 @@ COPY easyanimate/video_caption/requirements.txt /root/requirements-video_caption
|
||||
RUN pip install -r /root/requirements-video_caption.txt
|
||||
RUN rm /root/requirements-video_caption.txt
|
||||
|
||||
RUN apt-get update && apt-get install -y aria2
|
||||
RUN pip install -U http://eas-data.oss-cn-shanghai.aliyuncs.com/sdk/allspark-0.15-py2.py3-none-any.whl
|
||||
RUN pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
|
||||
RUN pip install came-pytorch deepspeed pytorch_lightning==1.9.4 func_timeout -i https://mirrors.aliyun.com/pypi/simple/
|
||||
@@ -59,9 +46,6 @@ RUN rm -rf /root/requirements.txt
|
||||
COPY easyanimate/video_caption/package_patches/easyocr_detection_patched.py /usr/local/lib/python3.10/dist-packages/easyocr/detection.py
|
||||
COPY easyanimate/video_caption/package_patches/vila_siglip_encoder_patched.py /usr/local/lib/python3.10/dist-packages/llava/model/multimodal_encoder/siglip_encoder.py
|
||||
|
||||
# Use the binary files provided by imageio-ffmpeg without installing the FFmpeg package (video-caption).
|
||||
RUN ln -s /usr/local/lib/python3.10/dist-packages/imageio_ffmpeg/binaries/ffmpeg-linux64-v4.2.2 /usr/local/bin/ffmpeg
|
||||
|
||||
ENV PYTHONUNBUFFERED 1
|
||||
ENV NVIDIA_DISABLE_REQUIRE 1
|
||||
|
||||
|
||||
@@ -27,16 +27,12 @@ allowing you to easily perform video data processing on [Aliyun PAI-DLC](https:/
|
||||
AliyunDSW or Docker is recommended to setup the environment, please refer to [Quick Start](../../README.md#quick-start).
|
||||
You can also refer to the image build process in the [Dockerfile](../../Dockerfile.ds) to configure the conda environment and other dependencies locally.
|
||||
|
||||
Since the video recaptioning depends on [llm-awq](https://github.com/mit-han-lab/llm-awq) for faster and memory efficient inference,
|
||||
the minimum GPU requirment should be RTX 3060 or A2 (CUDA Compute Capability >= 8.0).
|
||||
|
||||
```shell
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:
|
||||
asyanimate
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:easyanimate
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:asyanimate
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:easyanimate
|
||||
|
||||
# clone code
|
||||
git clone https://github.com/aigc-apps/EasyAnimate.git
|
||||
@@ -74,6 +70,8 @@ the video clips are obtained in `easyanimate/video_caption/datasets/panda_70m/vi
|
||||
Based on the videos obtained in the previous step, EasyAnimate provides a simple yet effective pipeline to filter out high-quality videos for recaptioning.
|
||||
The overall process is as follows:
|
||||
|
||||
- Scene transition filtering: Filter out videos with scene transition introduced by missing or superfluous splitting of PySceneDetect by calculating the semantic similarity
|
||||
accoss the beginning frame, the last frame, and the keyframes via [CLIP](https://github.com/openai/CLIP) or [DINOv2](https://github.com/facebookresearch/dinov2).
|
||||
- Aesthetic filtering: Filter out videos with poor content (blurry, dim, etc.) by calculating the average aesthetic score of uniformly sampled 4 frames via [aesthetic-predictor-v2-5](https://github.com/discus0434/aesthetic-predictor-v2-5).
|
||||
- Text filtering: Use [EasyOCR](https://github.com/JaidedAI/EasyOCR) to calculate the text area proportion of the middle frame to filter out videos with a large area of text.
|
||||
- Motion filtering: Calculate interframe optical flow differences to filter out videos that move too slowly or too quickly.
|
||||
@@ -83,23 +81,25 @@ After running
|
||||
```shell
|
||||
sh scripts/stage_2_video_filtering.sh
|
||||
```
|
||||
the aesthetic score, text score, and motion score of videos will be saved in the corresponding meta files in the folder `easyanimate/video_caption/datasets/panda_70m/videos_clips/`.
|
||||
the semantic consistency score, aesthetic score, text score, and motion score of videos will be saved
|
||||
in the corresponding meta files in the folder `easyanimate/video_caption/datasets/panda_70m/videos_clips/`.
|
||||
|
||||
> [!NOTE]
|
||||
> The computation of the aesthetic score depends on the [google/siglip-so400m-patch14-384 model](https://huggingface.co/google/siglip-so400m-patch14-384).
|
||||
> The computation of semantic consistency score depends on the [openai/clip-vit-large-patch14-336](https://huggingface.co/openai/clip-vit-large-patch14-336).
|
||||
Meanwhile, the aesthetic score depends on the [google/siglip-so400m-patch14-384 model](https://huggingface.co/google/siglip-so400m-patch14-384).
|
||||
Please run `HF_ENDPOINT=https://hf-mirror.com sh scripts/stage_2_video_filtering.sh` if you cannot access to huggingface.com.
|
||||
|
||||
|
||||
#### Video Recaptioning
|
||||
After obtaining the aboved high-quality filtered videos, EasyAnimate utilizes [VILA1.5](https://github.com/NVlabs/VILA) to perform video recaptioning.
|
||||
Subsequently, the recaptioning results are rewritten by LLMs to better meet with the requirements of video generation tasks.
|
||||
Finally, an advanced VideoCLIPXL model is developed to filter out video-caption pairs with poor alignment, resulting in the final training dataset.
|
||||
After obtaining the aboved high-quality filtered videos, EasyAnimate utilizes [InternVL2](https://internvl.readthedocs.io/en/latest/internvl2.0/introduction.html) to perform video recaptioning.
|
||||
Subsequently, the recaptioning results are rewritten by LLMs to better meet with the requirements of video generation tasks.
|
||||
Finally, an advanced [VideoCLIP-XL](https://arxiv.org/abs/2410.00741) model is used to filter out (video, long caption) pairs with poor alignment, resulting in the final training dataset.
|
||||
|
||||
Please download the video caption model from [VILA1.5](https://huggingface.co/collections/Efficient-Large-Model/vila-on-pre-training-for-visual-language-models-65d8022a3a52cd9bcd62698e) of the appropriate size based on the GPU memory of your machine.
|
||||
For A100 with 40G VRAM, you can download [VILA1.5-40b-AWQ](https://huggingface.co/Efficient-Large-Model/VILA1.5-40b-AWQ) by running
|
||||
Please download the video caption model from [InternVL2](https://huggingface.co/collections/OpenGVLab/internvl-20-667d3961ab5eb12c7ed1463e) of the appropriate size based on the GPU memory of your machine.
|
||||
For A100 with 40G VRAM, you can download [InternVL2-40B-AWQ](https://huggingface.co/OpenGVLab/InternVL2-40B-AWQ) by running
|
||||
```shell
|
||||
# Add HF_ENDPOINT=https://hf-mirror.com before the command if you cannot access to huggingface.com
|
||||
huggingface-cli download Efficient-Large-Model/VILA1.5-40b-AWQ --local-dir-use-symlinks False --local-dir /PATH/TO/VILA_MODEL
|
||||
huggingface-cli download OpenGVLab/InternVL2-40B-AWQ --local-dir-use-symlinks False --local-dir /PATH/TO/INTERNVL2_MODEL
|
||||
```
|
||||
|
||||
Optionally, you can prepare local LLMs to rewrite the recaption results.
|
||||
@@ -112,7 +112,7 @@ huggingface-cli download NousResearch/Meta-Llama-3-8B-Instruct --local-dir-use-s
|
||||
The entire workflow of video recaption is in the [stage_3_video_recaptioning.sh](./scripts/stage_3_video_recaptioning.sh).
|
||||
After running
|
||||
```shell
|
||||
VILA_MODEL_PATH=/PATH/TO/VILA_MODEL REWRITE_MODEL_PATH=/PATH/TO/REWRITE_MODEL sh scripts/stage_3_video_recaptioning.sh
|
||||
CAPTION_MODEL_PATH=/PATH/TO/INTERNVL2_MODEL REWRITE_MODEL_PATH=/PATH/TO/REWRITE_MODEL sh scripts/stage_3_video_recaptioning.sh
|
||||
```
|
||||
the final train file is obtained in `easyanimate/video_caption/datasets/panda_70m/videos_clips/meta_train_info.json`.
|
||||
|
||||
|
||||
@@ -24,15 +24,12 @@
|
||||
### 安装
|
||||
推荐使用阿里云 DSW 和 Docker 来安装环境，请参考 [快速开始](../../README_zh-CN.md#1-云使用-aliyundswdocker). 你也可以参考 [Dockerfile](../../Dockerfile.ds) 中的镜像构建流程在本地安装对应的 conda 环境和其余依赖。
|
||||
|
||||
为了提高推理速度和节省推理的显存，生成视频描述依赖于 [llm-awq](https://github.com/mit-han-lab/llm-awq)。因此，需要 RTX 3060 或者 A2 及以上的显卡 (CUDA Compute Capability >= 8.0)。
|
||||
|
||||
```shell
|
||||
# pull image
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:
|
||||
asyanimate
|
||||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:easyanimate
|
||||
|
||||
# enter image
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:asyanimate
|
||||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:easyanimate
|
||||
|
||||
# clone code
|
||||
git clone https://github.com/aigc-apps/EasyAnimate.git
|
||||
@@ -67,6 +64,7 @@ sh scripts/stage_1_video_splitting.sh
|
||||
#### 视频过滤
|
||||
基于上一步获得的视频，EasyAnimate 提供了一个简单而有效的流程来过滤出高质量的视频。总体流程如下：
|
||||
|
||||
- 场景跳变过滤：通过 [CLIP](https://github.com/openai/CLIP) 或者 [DINOv2](https://github.com/facebookresearch/dinov2) 来计算关键帧和首尾帧的语义相似度，从而过滤掉由于 PySceneDetect 缺失或多余分割引入的场景跳变的视频。
|
||||
- 美学过滤：通过 [aesthetic-predictor-v2-5](https://github.com/discus0434/aesthetic-predictor-v2-5) 计算均匀采样的 4 帧视频的平均美学分数，从而筛选出内容不佳（模糊、昏暗等）的视频。
|
||||
- 文本过滤：使用 [EasyOCR](https://github.com/JaidedAI/EasyOCR) 计算中间帧的文本区域比例，过滤掉含有大面积文本的视频。
|
||||
- 运动过滤：计算帧间光流差，过滤掉移动太慢或太快的视频。
|
||||
@@ -82,12 +80,12 @@ sh scripts/stage_2_video_filtering.sh
|
||||
请执行 `HF_ENDPOINT=https://hf-mirror.com sh scripts/stage_2_video_filtering.sh` 如果你无法访问 huggingface.com.
|
||||
|
||||
#### 视频描述
|
||||
在获得上述高质量的过滤视频后，EasyAnimate 利用 [VILA1.5](https://github.com/NVlabs/VILA) 来生成视频描述。随后，使用 LLMs 对生成的视频描述进行重写，以更好地满足视频生成任务的要求。最后，使用自研的 VideoCLIPXL 模型来过滤掉描述和视频内容不一致的数据，从而得到最终的训练数据集。
|
||||
在获得上述高质量的过滤视频后，EasyAnimate 利用 [InternVL2](https://internvl.readthedocs.io/en/latest/internvl2.0/introduction.html) 来生成视频描述。随后，使用 LLMs 对生成的视频描述进行重写，以更好地满足视频生成任务的要求。最后，使用自研的 [VideoCLIP-XL](https://arxiv.org/abs/2410.00741) 模型来过滤掉描述和视频内容不一致的数据，从而得到最终的训练数据集。
|
||||
|
||||
请根据机器的显存从 [VILA1.5](https://huggingface.co/collections/Efficient-Large-Model/vila-on-pre-training-for-visual-language-models-65d8022a3a52cd9bcd62698e) 下载合适大小的模型。对于 A100 40G，你可以执行下面的命令来下载 [VILA1.5-40b-AWQ](https://huggingface.co/Efficient-Large-Model/VILA1.5-40b-AWQ)
|
||||
请根据机器的显存从 [InternVL2](https://huggingface.co/collections/OpenGVLab/internvl-20-667d3961ab5eb12c7ed1463e) 下载合适大小的模型。对于 A100 40G，你可以执行下面的命令来下载 [InternVL2-40B-AWQ](https://huggingface.co/OpenGVLab/InternVL2-40B-AWQ)
|
||||
```shell
|
||||
# Add HF_ENDPOINT=https://hf-mirror.com before the command if you cannot access to huggingface.com
|
||||
huggingface-cli download Efficient-Large-Model/VILA1.5-40b-AWQ --local-dir-use-symlinks False --local-dir /PATH/TO/VILA_MODEL
|
||||
huggingface-cli download OpenGVLab/InternVL2-40B-AWQ --local-dir-use-symlinks False --local-dir /PATH/TO/INTERNVL2_MODEL
|
||||
```
|
||||
|
||||
你可以选择性地准备 LLMs 来改写上述视频描述的结果。例如，你执行下面的命令来下载 [Meta-Llama-3-8B-Instruct](https://huggingface.co/NousResearch/Meta-Llama-3-8B-Instruct)
|
||||
@@ -99,7 +97,7 @@ huggingface-cli download NousResearch/Meta-Llama-3-8B-Instruct --local-dir-use-s
|
||||
视频描述的完整流程在 [stage_3_video_recaptioning.sh](./scripts/stage_3_video_recaptioning.sh).
|
||||
执行
|
||||
```shell
|
||||
VILA_MODEL_PATH=/PATH/TO/VILA_MODEL REWRITE_MODEL_PATH=/PATH/TO/REWRITE_MODEL sh scripts/stage_3_video_recaptioning.sh
|
||||
CAPTION_MODEL_PATH=/PATH/TO/INTERNVL2_MODEL REWRITE_MODEL_PATH=/PATH/TO/REWRITE_MODEL sh scripts/stage_3_video_recaptioning.sh
|
||||
```
|
||||
后，最后的训练文件会保存在 `easyanimate/video_caption/datasets/panda_70m/videos_clips/meta_train_info.json`。
|
||||
|
||||
|
||||
@@ -7,9 +7,10 @@ import torch
|
||||
from natsort import index_natsorted
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer
|
||||
from utils.logger import logger
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
from utils.logger import logger
|
||||
|
||||
|
||||
def extract_output(s, prefix='"rewritten description": '):
|
||||
"""Customize the function according to the prompt."""
|
||||
@@ -40,7 +41,7 @@ def parse_args():
|
||||
parser.add_argument(
|
||||
"--video_path_column",
|
||||
type=str,
|
||||
default=None,
|
||||
default=None, # In the beautiful prompt case, it is not necessary.
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -175,6 +176,7 @@ def main():
|
||||
batch_video_path, batch_output = zip(*batch_result)
|
||||
|
||||
result_dict[args.video_path_column].extend(batch_video_path)
|
||||
result_dict[args.caption_column].extend(batch_output)
|
||||
else:
|
||||
for output in batch_output:
|
||||
if output is not None:
|
||||
@@ -183,41 +185,25 @@ def main():
|
||||
result_dict[args.caption_column].extend(batch_result)
|
||||
|
||||
# Save the metadata every args.saved_freq.
|
||||
if i != 0 and ((i // args.batch_size) % args.saved_freq) == 0:
|
||||
if (i // args.batch_size) % args.saved_freq == 0 or (i + 1) * args.batch_size >= len(sampled_frame_caption_list):
|
||||
if len(result_dict[args.caption_column]) > 0:
|
||||
result_df = pd.DataFrame(result_dict)
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = True if not os.path.exists(args.saved_path) else False
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
elif args.saved_path.endswith(".json"):
|
||||
# Append is not supported.
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records")
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", indent=4, force_ascii=False)
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
|
||||
result_dict = {args.caption_column: []}
|
||||
if args.video_path_column is not None:
|
||||
result_dict = {args.video_path_column: [], args.caption_column: []}
|
||||
|
||||
if len(result_dict[args.caption_column]) > 0:
|
||||
result_df = pd.DataFrame(result_dict)
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = True if not os.path.exists(args.saved_path) else False
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a")
|
||||
elif args.saved_path.endswith(".json"):
|
||||
# Append is not supported.
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records")
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", indent=4, force_ascii=False)
|
||||
logger.info(f"Save the final result to {args.saved_path}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -11,6 +11,7 @@ import pandas as pd
|
||||
from joblib import Parallel, delayed
|
||||
from natsort import natsorted
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
|
||||
@@ -98,27 +99,35 @@ def parse_args():
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).")
|
||||
parser.add_argument("--saved_freq", type=int, default=100, help="The frequency to save the output results.")
|
||||
parser.add_argument("--saved_freq", type=int, default=1, help="The frequency to save the output results.")
|
||||
parser.add_argument("--n_jobs", type=int, default=1, help="The number of concurrent processes.")
|
||||
|
||||
parser.add_argument(
|
||||
"--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl).")
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
"--aesthetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score", type=float, default=4.0, help="The asethetic score threshold.")
|
||||
parser.add_argument("--min_aesthetic_score", type=float, default=4.0, help="The aesthetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
"--aesthetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score_siglip", type=float, default=4.0, help="The asethetic score (SigLIP) threshold.")
|
||||
parser.add_argument("--min_aesthetic_score_siglip", type=float, default=4.0, help="The aesthetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_text_score", type=float, default=0.02, help="The text threshold.")
|
||||
parser.add_argument(
|
||||
"--semantic_consistency_score_metadata_path",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the semantic consistency metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min_semantic_consistency_score", type=float, default=0.80, help="The semantic consistency score threshold."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -153,12 +162,13 @@ def main():
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
asethetic_score_metadata_path=args.asethetic_score_metadata_path,
|
||||
min_asethetic_score=args.min_asethetic_score,
|
||||
asethetic_score_siglip_metadata_path=args.asethetic_score_siglip_metadata_path,
|
||||
min_asethetic_score_siglip=args.min_asethetic_score_siglip,
|
||||
aesthetic_score_metadata_path=args.aesthetic_score_metadata_path,
|
||||
min_aesthetic_score=args.min_aesthetic_score,
|
||||
aesthetic_score_siglip_metadata_path=args.aesthetic_score_siglip_metadata_path,
|
||||
min_aesthetic_score_siglip=args.min_aesthetic_score_siglip,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
video_path_column=args.video_path_column
|
||||
)
|
||||
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list]
|
||||
# Sorting to guarantee the same result for each process.
|
||||
@@ -173,13 +183,18 @@ def main():
|
||||
continue
|
||||
|
||||
result_df = pd.DataFrame(result_list)
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,225 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from accelerate import PartialState
|
||||
from accelerate.utils import gather_object
|
||||
from natsort import natsorted
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoImageProcessor, AutoModel
|
||||
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
from utils.video_dataset import VideoDataset, collate_fn
|
||||
from utils.video_utils import ALL_FRAME_SAMPLE_METHODS
|
||||
|
||||
|
||||
ALL_MODEL_NAME = [
|
||||
"dinov2-small",
|
||||
"dinov2-base",
|
||||
"dinov2-large",
|
||||
"clip-vit-large-patch14",
|
||||
"clip-vit-base-patch32",
|
||||
"clip-vit-large-patch14-336",
|
||||
]
|
||||
|
||||
|
||||
def init_model(model_name, device):
|
||||
processor = AutoImageProcessor.from_pretrained(model_name)
|
||||
model = AutoModel.from_pretrained(model_name).to(device)
|
||||
return processor, model
|
||||
|
||||
|
||||
def compute_adjacent_similarity(frame_features):
|
||||
frame_features /= frame_features.norm(dim=-1, keepdim=True)
|
||||
roll_frame_features = torch.roll(frame_features, shifts=-1, dims=0)
|
||||
similarity_matrix = frame_features.squeeze(dim=1).cpu().numpy() @ roll_frame_features.squeeze(dim=1).cpu().numpy().T
|
||||
|
||||
return np.diag(similarity_matrix).tolist()[:-1]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Compute the semantic consistency score across frames.")
|
||||
parser.add_argument(
|
||||
"--video_metadata_path", type=str, required=True, help="The path to the video dataset metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_path_column",
|
||||
type=str,
|
||||
default="video_path",
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument("--video_folder", type=str, default="", help="The video folder.")
|
||||
parser.add_argument(
|
||||
"--model_path", type=str, default="openai/clip-vit-large-patch14-336", help="The path to the DINO/CLIP model."
|
||||
)
|
||||
parser.add_argument("--frame_sample_method", type=str, choices=ALL_FRAME_SAMPLE_METHODS, default="keyframe+first")
|
||||
parser.add_argument("--num_sampled_frames", type=int, default=1, help="The number of sampled frames.")
|
||||
parser.add_argument("--sample_stride", type=int, default=None, help="The stride between two sampled frames.")
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="The batch size for the video dataset.")
|
||||
parser.add_argument("--num_workers", type=int, default=1, help="The number of workers for the video dataset.")
|
||||
parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).")
|
||||
parser.add_argument("--saved_freq", type=int, default=1, help="The frequency to save the output results.")
|
||||
|
||||
parser.add_argument("--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl).")
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--aesthetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_aesthetic_score", type=float, default=4.0, help="The aesthetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--aesthetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_aesthetic_score_siglip", type=float, default=4.0, help="The aesthetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_text_score", type=float, default=0.02, help="The text threshold.")
|
||||
parser.add_argument(
|
||||
"--motion_score_metadata_path", type=str, default=None, help="The path to the video motion score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The motion threshold.")
|
||||
parser.add_argument("--max_motion_score", type=float, default=999999, help="The maximum motion threshold.")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.video_metadata_path.endswith(".csv"):
|
||||
video_metadata_df = pd.read_csv(args.video_metadata_path)
|
||||
elif args.video_metadata_path.endswith(".jsonl"):
|
||||
video_metadata_df = pd.read_json(args.video_metadata_path, lines=True)
|
||||
else:
|
||||
raise ValueError("The video_metadata_path must end with .csv or .jsonl.")
|
||||
video_path_list = video_metadata_df[args.video_path_column].tolist()
|
||||
|
||||
if not (args.saved_path.endswith(".csv") or args.saved_path.endswith(".jsonl")):
|
||||
raise ValueError("The saved_path must end with .csv or .jsonl.")
|
||||
|
||||
if os.path.exists(args.saved_path):
|
||||
if args.saved_path.endswith(".csv"):
|
||||
saved_metadata_df = pd.read_csv(args.saved_path)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
saved_metadata_df = pd.read_json(args.saved_path, lines=True)
|
||||
saved_video_path_list = saved_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = list(set(video_path_list).difference(set(saved_video_path_list)))
|
||||
logger.info(f"Resume from {args.saved_path}: {len(saved_video_path_list)} processed and {len(video_path_list)} to be processed.")
|
||||
|
||||
video_path_list = filter(
|
||||
video_path_list,
|
||||
basic_metadata_path=args.basic_metadata_path,
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
aesthetic_score_metadata_path=args.aesthetic_score_metadata_path,
|
||||
min_aesthetic_score=args.min_aesthetic_score,
|
||||
aesthetic_score_siglip_metadata_path=args.aesthetic_score_siglip_metadata_path,
|
||||
min_aesthetic_score_siglip=args.min_aesthetic_score_siglip,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
min_motion_score=args.min_motion_score,
|
||||
max_motion_score=args.max_motion_score,
|
||||
video_path_column=args.video_path_column
|
||||
)
|
||||
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list]
|
||||
# Sorting to guarantee the same result for each process.
|
||||
video_path_list = natsorted(video_path_list)
|
||||
|
||||
if not any(name in args.model_path for name in ALL_MODEL_NAME):
|
||||
raise ValueError(f"The model_path should be among the following list: {ALL_MODEL_NAME}.")
|
||||
|
||||
state = PartialState()
|
||||
if state.is_main_process:
|
||||
# Check if the model is downloaded in the main process.
|
||||
processor, model = init_model(args.model_path, "cpu")
|
||||
state.wait_for_everyone()
|
||||
processor, model = init_model(args.model_path, state.device)
|
||||
|
||||
index = len(video_path_list) - len(video_path_list) % state.num_processes
|
||||
# Avoid the NCCL timeout in the final gather operation.
|
||||
logger.warning(
|
||||
f"Drop the last {len(video_path_list) % state.num_processes} videos "
|
||||
"to ensure each process handles the same number of videos."
|
||||
)
|
||||
video_path_list = video_path_list[:index]
|
||||
logger.info(f"{len(video_path_list)} videos are to be processed.")
|
||||
|
||||
result_dict = {
|
||||
args.video_path_column: [],
|
||||
"similarity_cross_frame": [],
|
||||
"similarity_mean": [],
|
||||
"sample_frame_idx": [],
|
||||
}
|
||||
with state.split_between_processes(video_path_list) as splitted_video_path_list:
|
||||
video_dataset = VideoDataset(
|
||||
dataset_inputs={args.video_path_column: splitted_video_path_list},
|
||||
video_folder=args.video_folder,
|
||||
sample_method=args.frame_sample_method,
|
||||
num_sampled_frames=args.num_sampled_frames,
|
||||
sample_stride=args.sample_stride,
|
||||
)
|
||||
video_loader = DataLoader(video_dataset, batch_size=args.batch_size, num_workers=args.num_workers, collate_fn=collate_fn)
|
||||
|
||||
for idx, batch in enumerate(tqdm(video_loader)):
|
||||
if len(batch) > 0:
|
||||
batch_video_path = batch["path"]
|
||||
batch_frame = batch["sampled_frame"]
|
||||
frame_num_list = [len(video_frames) for video_frames in batch_frame]
|
||||
# [B, T, H, W, C] => [(B * T), H, W, C]
|
||||
reshaped_batch_frame = [frame for video_frames in batch_frame for frame in video_frames]
|
||||
with torch.no_grad():
|
||||
inputs = processor(images=reshaped_batch_frame, return_tensors="pt").to(state.device)
|
||||
if "dino" in args.model_path.lower():
|
||||
frame_features = model(**inputs).last_hidden_state.mean(dim=1)
|
||||
else: # CLIP
|
||||
frame_features = model.get_image_features(**inputs)
|
||||
|
||||
# Each video may have a different number of sampled frames.
|
||||
# Map the flattened frame features back to their original shape.
|
||||
batch_frame_features = torch.split(frame_features, frame_num_list)
|
||||
batch_simi_cross_frame = [compute_adjacent_similarity(frame_features) for frame_features in batch_frame_features]
|
||||
batch_similarity_mean = [
|
||||
sum(simi_cross_frame) / len(simi_cross_frame) for simi_cross_frame in batch_simi_cross_frame
|
||||
]
|
||||
|
||||
if args.video_folder == "":
|
||||
saved_video_path_list = batch_video_path
|
||||
else:
|
||||
saved_video_path_list = [os.path.relpath(video_path, args.video_folder) for video_path in batch_video_path]
|
||||
result_dict[args.video_path_column].extend(saved_video_path_list)
|
||||
result_dict["similarity_cross_frame"].extend(batch_simi_cross_frame)
|
||||
result_dict["similarity_mean"].extend(batch_similarity_mean)
|
||||
result_dict["sample_frame_idx"].extend(batch["sampled_frame_idx"])
|
||||
|
||||
# Save the metadata in the main process every saved_freq.
|
||||
if (idx != 0) and (idx % args.saved_freq == 0 or idx == len(video_loader) - 1):
|
||||
state.wait_for_everyone()
|
||||
gathered_result_dict = {k: gather_object(v) for k, v in result_dict.items()}
|
||||
if state.is_main_process and len(gathered_result_dict[args.video_path_column]) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_dict)
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
for k in result_dict.keys():
|
||||
result_dict[k] = []
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -8,8 +8,9 @@ import pandas as pd
|
||||
from accelerate import PartialState
|
||||
from accelerate.utils import gather_object
|
||||
from natsort import natsorted
|
||||
from torchvision.datasets.utils import download_url
|
||||
from tqdm import tqdm
|
||||
from torchvision.datasets.utils import download_url
|
||||
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
from utils.video_utils import extract_frames
|
||||
@@ -98,26 +99,35 @@ def parse_args():
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).")
|
||||
parser.add_argument("--saved_freq", type=int, default=100, help="The frequency to save the output results.")
|
||||
parser.add_argument("--saved_freq", type=int, default=1, help="The frequency to save the output results.")
|
||||
|
||||
parser.add_argument(
|
||||
"--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl).")
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
"--aesthetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score", type=float, default=4.0, help="The asethetic score threshold.")
|
||||
parser.add_argument("--min_aesthetic_score", type=float, default=4.0, help="The aesthetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
"--aesthetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score_siglip", type=float, default=4.0, help="The asethetic score (SigLIP) threshold.")
|
||||
parser.add_argument("--min_aesthetic_score_siglip", type=float, default=4.0, help="The aesthetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--motion_score_metadata_path", type=str, default=None, help="The path to the video motion score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The motion threshold.")
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The minimum motion threshold.")
|
||||
parser.add_argument("--max_motion_score", type=float, default=999999, help="The maximum motion threshold.")
|
||||
parser.add_argument(
|
||||
"--semantic_consistency_score_metadata_path",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the semantic consistency metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min_semantic_consistency_score", type=float, default=0.80, help="The semantic consistency score threshold."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -152,12 +162,16 @@ def main():
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
asethetic_score_metadata_path=args.asethetic_score_metadata_path,
|
||||
min_asethetic_score=args.min_asethetic_score,
|
||||
asethetic_score_siglip_metadata_path=args.asethetic_score_siglip_metadata_path,
|
||||
min_asethetic_score_siglip=args.min_asethetic_score_siglip,
|
||||
aesthetic_score_metadata_path=args.aesthetic_score_metadata_path,
|
||||
min_aesthetic_score=args.min_aesthetic_score,
|
||||
aesthetic_score_siglip_metadata_path=args.aesthetic_score_siglip_metadata_path,
|
||||
min_aesthetic_score_siglip=args.min_aesthetic_score_siglip,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
min_motion_score=args.min_motion_score,
|
||||
max_motion_score=args.max_motion_score,
|
||||
semantic_consistency_score_metadata_path=args.semantic_consistency_score_metadata_path,
|
||||
min_semantic_consistency_score=args.min_semantic_consistency_score,
|
||||
video_path_column=args.video_path_column
|
||||
)
|
||||
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list]
|
||||
# Sorting to guarantee the same result for each process.
|
||||
@@ -172,7 +186,10 @@ def main():
|
||||
|
||||
index = len(video_path_list) - len(video_path_list) % state.num_processes
|
||||
# Avoid the NCCL timeout in the final gather operation.
|
||||
logger.info(f"Drop {len(video_path_list) % state.num_processes} videos to ensure each process handles the same number of videos.")
|
||||
logger.info(
|
||||
f"Drop the last {len(video_path_list) % state.num_processes} videos to "
|
||||
"ensure each process handles the same number of videos."
|
||||
)
|
||||
video_path_list = video_path_list[:index]
|
||||
logger.info(f"{len(video_path_list)} videos are to be processed.")
|
||||
|
||||
@@ -184,30 +201,24 @@ def main():
|
||||
result_list.append(video_meta_info)
|
||||
except Exception as e:
|
||||
logger.warning(f"Compute text score for video {video_path} with error: {e}.")
|
||||
if i != 0 and i % args.saved_freq == 0:
|
||||
if i % args.saved_freq == 0 or i == len(splitted_video_path_list) - 1:
|
||||
state.wait_for_everyone()
|
||||
gathered_result_list = gather_object(result_list)
|
||||
if state.is_main_process and len(gathered_result_list) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_list)
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
result_list = []
|
||||
|
||||
state.wait_for_everyone()
|
||||
gathered_result_list = gather_object(result_list)
|
||||
if state.is_main_process and len(gathered_result_list) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_list)
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
logger.info(f"Save the final result to {args.saved_path}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -2,13 +2,15 @@ import argparse
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
import utils.image_evaluator as image_evaluator
|
||||
import utils.video_evaluator as video_evaluator
|
||||
from accelerate import PartialState
|
||||
from accelerate.utils import gather_object
|
||||
from natsort import index_natsorted
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
import utils.image_evaluator as image_evaluator
|
||||
import utils.video_evaluator as video_evaluator
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
from utils.video_dataset import VideoDataset, collate_fn
|
||||
|
||||
@@ -25,41 +27,43 @@ def parse_args():
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument("--video_folder", type=str, default="", help="The video folder.")
|
||||
parser.add_argument(
|
||||
"--caption_column",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The column contains the caption.",
|
||||
)
|
||||
parser.add_argument("--caption_column", type=str, default=None, help="The column contains the caption.")
|
||||
parser.add_argument(
|
||||
"--frame_sample_method",
|
||||
type=str,
|
||||
choices=["mid", "uniform", "image"],
|
||||
default="uniform",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_sampled_frames",
|
||||
type=int,
|
||||
default=8,
|
||||
help="num_sampled_frames",
|
||||
)
|
||||
parser.add_argument("--num_sampled_frames", type=int, default=8, help="The number of sampled frames.")
|
||||
parser.add_argument("--metrics", nargs="+", type=str, required=True, help="The evaluation metric(s) for generated images.")
|
||||
parser.add_argument(
|
||||
"--batch_size",
|
||||
type=int,
|
||||
default=10,
|
||||
required=False,
|
||||
help="The batch size for the video dataset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_workers",
|
||||
type=int,
|
||||
default=4,
|
||||
required=False,
|
||||
help="The number of workers for the video dataset.",
|
||||
)
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="The batch size for the video dataset.")
|
||||
parser.add_argument("--num_workers", type=int, default=1, help="The number of workers for the video dataset.")
|
||||
parser.add_argument("--saved_path", type=str, required=True, help="The save path to the output results (csv/jsonl).")
|
||||
parser.add_argument("--saved_freq", type=int, default=1000, help="The frequency to save the output results.")
|
||||
parser.add_argument("--saved_freq", type=int, default=1, help="The frequency to save the output results.")
|
||||
|
||||
parser.add_argument("--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl).")
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_text_score", type=float, default=0.02, help="The text threshold.")
|
||||
parser.add_argument(
|
||||
"--motion_score_metadata_path", type=str, default=None, help="The path to the video motion score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The minimum motion threshold.")
|
||||
parser.add_argument("--max_motion_score", type=float, default=999999, help="The maximum motion threshold.")
|
||||
parser.add_argument(
|
||||
"--semantic_consistency_score_metadata_path",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the semantic consistency metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min_semantic_consistency_score", type=float, default=0.80, help="The semantic consistency score threshold."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -95,6 +99,24 @@ def main():
|
||||
video_metadata_df = video_metadata_df[[args.video_path_column, args.caption_column + "_x"]]
|
||||
video_metadata_df.rename(columns={args.caption_column + "_x": args.caption_column}, inplace=True)
|
||||
logger.info(f"Resume from {args.saved_path}: {len(saved_metadata_df)} processed and {len(video_metadata_df)} to be processed.")
|
||||
|
||||
video_path_list = video_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = filter(
|
||||
video_path_list,
|
||||
basic_metadata_path=args.basic_metadata_path,
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
min_motion_score=args.min_motion_score,
|
||||
max_motion_score=args.max_motion_score,
|
||||
semantic_consistency_score_metadata_path=args.semantic_consistency_score_metadata_path,
|
||||
min_semantic_consistency_score=args.min_semantic_consistency_score,
|
||||
video_path_column=args.video_path_column
|
||||
)
|
||||
video_metadata_df = video_metadata_df[video_metadata_df[args.video_path_column].isin(video_path_list)]
|
||||
|
||||
state = PartialState()
|
||||
metric_fns = []
|
||||
@@ -126,7 +148,10 @@ def main():
|
||||
|
||||
index = len(video_metadata_df) - len(video_metadata_df) % state.num_processes
|
||||
# Avoid the NCCL timeout in the final gather operation.
|
||||
logger.info(f"Drop {len(video_metadata_df) % state.num_processes} videos to ensure each process handles the same number of videos.")
|
||||
logger.info(
|
||||
f"Drop the last {len(video_metadata_df) % state.num_processes} videos "
|
||||
"to ensure each process handles the same number of videos."
|
||||
)
|
||||
video_metadata_df = video_metadata_df.iloc[:index]
|
||||
logger.info(f"{len(video_metadata_df)} videos are to be processed.")
|
||||
|
||||
@@ -169,32 +194,25 @@ def main():
|
||||
result_dict[args.video_path_column].extend(saved_video_path_list)
|
||||
|
||||
# Save the metadata in the main process every saved_freq.
|
||||
if (idx != 0) and (idx % args.saved_freq == 0):
|
||||
if (idx % args.saved_freq) == 0 or idx == len(video_loader) - 1:
|
||||
state.wait_for_everyone()
|
||||
gathered_result_dict = {k: gather_object(v) for k, v in result_dict.items()}
|
||||
if state.is_main_process and len(gathered_result_dict[args.video_path_column]) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_dict)
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
for k in result_dict.keys():
|
||||
result_dict[k] = []
|
||||
|
||||
# Wait for all processes to finish and gather the final result.
|
||||
state.wait_for_everyone()
|
||||
gathered_result_dict = {k: gather_object(v) for k, v in result_dict.items()}
|
||||
# Save the metadata in the main process.
|
||||
if state.is_main_process and len(gathered_result_dict[args.video_path_column]) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_dict)
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
logger.info(f"Save the final result to {args.saved_path}.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -26,13 +26,13 @@ def parse_args():
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
"--aesthetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score", type=float, default=4.0, help="The asethetic score threshold.")
|
||||
parser.add_argument("--min_aesthetic_score", type=float, default=4.0, help="The aesthetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality (SigLIP) metadata (csv/jsonl)."
|
||||
"--aesthetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality (SigLIP) metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score_siglip", type=float, default=4.0, help="The asethetic score (SigLIP) threshold.")
|
||||
parser.add_argument("--min_aesthetic_score_siglip", type=float, default=4.0, help="The aesthetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
@@ -62,10 +62,10 @@ def main():
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
asethetic_score_metadata_path=args.asethetic_score_metadata_path,
|
||||
min_asethetic_score=args.min_asethetic_score,
|
||||
asethetic_score_siglip_metadata_path=args.asethetic_score_siglip_metadata_path,
|
||||
min_asethetic_score_siglip=args.min_asethetic_score_siglip,
|
||||
aesthetic_score_metadata_path=args.aesthetic_score_metadata_path,
|
||||
min_aesthetic_score=args.min_aesthetic_score,
|
||||
aesthetic_score_siglip_metadata_path=args.aesthetic_score_siglip_metadata_path,
|
||||
min_aesthetic_score_siglip=args.min_aesthetic_score_siglip,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
from natsort import natsorted
|
||||
from transformers import AutoTokenizer
|
||||
from tqdm import tqdm
|
||||
from torch.utils.data import DataLoader
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
from utils.video_dataset import VideoDataset, collate_fn
|
||||
|
||||
|
||||
def recaption_batch_video(llm, batch_video_frames, prompt, sampling_params):
|
||||
inputs = [
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"image": video_frames
|
||||
},
|
||||
}
|
||||
for video_frames in batch_video_frames
|
||||
]
|
||||
|
||||
outputs = llm.generate(inputs, sampling_params=sampling_params)
|
||||
|
||||
batch_output = []
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
batch_output.append(generated_text)
|
||||
|
||||
return batch_output
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Recaption videos with InternVL2.")
|
||||
parser.add_argument(
|
||||
"--video_metadata_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the video dataset metadata (csv/jsonl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_path_column",
|
||||
type=str,
|
||||
default="video_path",
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--caption_column",
|
||||
type=str,
|
||||
default="caption",
|
||||
help="The column contains the caption.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_folder", type=str, default="", help="The video folder."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size",
|
||||
type=int,
|
||||
default=1,
|
||||
required=False,
|
||||
help="The batch size for vllm inference. Adjust according to the number of GPUs to maximize inference throughput.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_workers",
|
||||
type=int,
|
||||
default=4,
|
||||
required=False,
|
||||
help="The number of workers for the video dataset.",
|
||||
)
|
||||
parser.add_argument("--input_prompt", type=str, default="Describe this video in detail. Don\'t repeat.")
|
||||
parser.add_argument(
|
||||
"--model_path", type=str, default="OpenGVLab/InternVL2-40B-AWQ"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame_sample_method",
|
||||
type=str,
|
||||
choices=["mid", "uniform", "image"],
|
||||
default="uniform",
|
||||
)
|
||||
parser.add_argument("--num_sampled_frames", type=int, default=8)
|
||||
parser.add_argument(
|
||||
"--saved_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The save path to the output results (csv/jsonl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--saved_freq",
|
||||
type=int,
|
||||
default=1,
|
||||
help="The frequency to save the output results.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--aesthetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_aesthetic_score", type=float, default=4.0, help="The aesthetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--aesthetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_aesthetic_score_siglip", type=float, default=4.0, help="The aesthetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_text_score", type=float, default=0.02, help="The text threshold.")
|
||||
parser.add_argument(
|
||||
"--motion_score_metadata_path", type=str, default=None, help="The path to the video motion score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The motion threshold.")
|
||||
parser.add_argument("--max_motion_score", type=float, default=999999, help="The maximum motion threshold.")
|
||||
parser.add_argument(
|
||||
"--semantic_consistency_score_metadata_path",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the semantic consistency metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min_semantic_consistency_score", type=float, default=0.80, help="The semantic consistency score threshold."
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.video_metadata_path.endswith(".csv"):
|
||||
video_metadata_df = pd.read_csv(args.video_metadata_path)
|
||||
elif args.video_metadata_path.endswith(".jsonl"):
|
||||
video_metadata_df = pd.read_json(args.video_metadata_path, lines=True)
|
||||
else:
|
||||
raise ValueError("The video_metadata_path must end with .csv or .jsonl.")
|
||||
video_path_list = video_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = [os.path.basename(video_path) for video_path in video_path_list]
|
||||
|
||||
if not (args.saved_path.endswith(".csv") or args.saved_path.endswith(".jsonl")):
|
||||
raise ValueError("The saved_path must end with .csv or .jsonl.")
|
||||
|
||||
if os.path.exists(args.saved_path):
|
||||
if args.saved_path.endswith(".csv"):
|
||||
saved_metadata_df = pd.read_csv(args.saved_path)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
saved_metadata_df = pd.read_json(args.saved_path, lines=True)
|
||||
saved_video_path_list = saved_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = list(set(video_path_list).difference(set(saved_video_path_list)))
|
||||
logger.info(
|
||||
f"Resume from {args.saved_path}: {len(saved_video_path_list)} processed and {len(video_path_list)} to be processed."
|
||||
)
|
||||
|
||||
video_path_list = filter(
|
||||
video_path_list,
|
||||
basic_metadata_path=args.basic_metadata_path,
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
aesthetic_score_metadata_path=args.aesthetic_score_metadata_path,
|
||||
min_aesthetic_score=args.min_aesthetic_score,
|
||||
aesthetic_score_siglip_metadata_path=args.aesthetic_score_siglip_metadata_path,
|
||||
min_aesthetic_score_siglip=args.min_aesthetic_score_siglip,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
min_motion_score=args.min_motion_score,
|
||||
)
|
||||
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list]
|
||||
# Sorting to guarantee the same result for each process.
|
||||
video_path_list = natsorted(video_path_list)
|
||||
|
||||
video_dataset = VideoDataset(
|
||||
dataset_inputs={args.video_path_column: video_path_list},
|
||||
video_folder=args.video_folder,
|
||||
sample_method=args.frame_sample_method,
|
||||
num_sampled_frames=args.num_sampled_frames
|
||||
)
|
||||
video_loader = DataLoader(video_dataset, batch_size=args.batch_size, num_workers=args.num_workers, collate_fn=collate_fn)
|
||||
|
||||
# Initialize the vllm inference pipeline.
|
||||
CUDA_VISIBLE_DEVICES = os.getenv("CUDA_VISIBLE_DEVICES", None)
|
||||
tensor_parallel_size = torch.cuda.device_count() if CUDA_VISIBLE_DEVICES is None else len(CUDA_VISIBLE_DEVICES.split(","))
|
||||
logger.info(f"Automatically set tensor_parallel_size={tensor_parallel_size} based on the available devices.")
|
||||
|
||||
llm = LLM(
|
||||
model=args.model_path,
|
||||
trust_remote_code=True,
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"image": args.num_sampled_frames},
|
||||
gpu_memory_utilization=0.9,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
quantization="awq",
|
||||
dtype="float16",
|
||||
mm_processor_kwargs={"max_dynamic_patch": 1}
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
|
||||
placeholders = "".join(f"Frame{i}: <image>\n" for i in range(1, args.num_sampled_frames + 1))
|
||||
messages = [{'role': 'user', 'content': f"{placeholders}{args.input_prompt}"}]
|
||||
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
# Stop tokens for InternVL
|
||||
# models variants may have different stop tokens
|
||||
# please refer to the model card for the correct "stop words":
|
||||
# https://huggingface.co/OpenGVLab/InternVL2-2B#service
|
||||
stop_tokens = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|end|>"]
|
||||
stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens]
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=512, stop_token_ids=stop_token_ids)
|
||||
|
||||
result_dict = {args.video_path_column: [], args.caption_column: []}
|
||||
for idx, batch in enumerate(tqdm(video_loader)):
|
||||
if len(batch) > 0:
|
||||
batch_video_path = batch["path"]
|
||||
batch_frame = batch["sampled_frame"] # [batch_size, num_sampled_frames, H, W, C]
|
||||
batch_caption = recaption_batch_video(llm, batch_frame, prompt, sampling_params)
|
||||
|
||||
if args.video_folder == "":
|
||||
saved_video_path_list = batch_video_path
|
||||
else:
|
||||
saved_video_path_list = [os.path.relpath(video_path, args.video_folder) for video_path in batch_video_path]
|
||||
result_dict[args.video_path_column].extend(saved_video_path_list)
|
||||
result_dict["caption"].extend(batch_caption)
|
||||
|
||||
if idx % args.saved_freq == 0 or idx == len(video_loader) - 1:
|
||||
result_df = pd.DataFrame(result_dict)
|
||||
|
||||
# Append is not supported (oss).
|
||||
if args.saved_path.endswith(".csv"):
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_csv(args.saved_path)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df = result_df.iloc[natsorted(result_df.index, key=lambda x: result_df.loc[x, args.video_path_column])]
|
||||
result_df.to_csv(args.saved_path, index=False)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
if os.path.exists(args.saved_path):
|
||||
saved_df = pd.read_json(args.saved_path, orient="records", lines=True)
|
||||
result_df = pd.concat([saved_df, result_df], ignore_index=True)
|
||||
result_df = result_df.iloc[natsorted(result_df.index, key=lambda x: result_df.loc[x, args.video_path_column])]
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
result_dict = {args.video_path_column: [], args.caption_column: []}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -2,8 +2,9 @@ Please rewrite the video description to be useful for AI to re-generate the vide
|
||||
1. Do not start with something similar to 'The video/scene/frame shows' or "In this video/scene/frame".
|
||||
2. Remove the subjective content deviates from describing the visual content of the video. For instance, a sentence like "It gives a feeling of ease and tranquility and makes people feel comfortable" is considered subjective.
|
||||
3. Remove the non-existent description that does not in the visual content of the video, For instance, a sentence like "There is no visible detail that could be used to identify the individual beyond what is shown." is considered as the non-existent description.
|
||||
4. Here are some examples of good descriptions: 1) A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about. 2) A large orange octopus is seen resting on the bottom of the ocean floor, blending in with the sandy and rocky terrain. Its tentacles are spread out around its body, and its eyes are closed. The octopus is unaware of a king crab that is crawling towards it from behind a rock, its claws raised and ready to attack. The crab is brown and spiny, with long legs and antennae. The scene is captured from a wide angle, showing the vastness and depth of the ocean. The water is clear and blue, with rays of sunlight filtering through. The shot is sharp and crisp, with a high dynamic range. The octopus and the crab are in focus, while the background is slightly blurred, creating a depth of field effect.
|
||||
5. Output with the following json format:
|
||||
4. The rewritten description should include the main subject (person, object, animal, or none) actions and their attributes or status sequence, the background (the objects, location, weather, and time), the view shot, camera movement and the video style.
|
||||
5. Here are some examples of good descriptions: 1) A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about. 2) A large orange octopus is seen resting on the bottom of the ocean floor, blending in with the sandy and rocky terrain. Its tentacles are spread out around its body, and its eyes are closed. The octopus is unaware of a king crab that is crawling towards it from behind a rock, its claws raised and ready to attack. The crab is brown and spiny, with long legs and antennae. The scene is captured from a wide angle, showing the vastness and depth of the ocean. The water is clear and blue, with rays of sunlight filtering through. The shot is sharp and crisp, with a high dynamic range. The octopus and the crab are in focus, while the background is slightly blurred, creating a depth of field effect.
|
||||
6. Output with the following json format:
|
||||
{"rewritten description": "your rewritten description here"}
|
||||
|
||||
Here is the video description:
|
||||
@@ -4,6 +4,4 @@ git+https://github.com/openai/CLIP.git
|
||||
natsort
|
||||
joblib
|
||||
scenedetect
|
||||
av
|
||||
# https://github.com/NVlabs/VILA/issues/78#issuecomment-2195568292
|
||||
numpy<2.0.0
|
||||
av
|
||||
@@ -1,7 +1,10 @@
|
||||
META_FILE_PATH="datasets/panda_70m/videos_clips/data/meta_file_info.jsonl"
|
||||
VIDEO_FOLDER="datasets/panda_70m/videos_clips/data/"
|
||||
VIDEO_QUALITY_SAVED_PATH="datasets/panda_70m/videos_clips/meta_quality_info_siglip.jsonl"
|
||||
MIN_ASETHETIC_SCORE_SIGLIP=4.0
|
||||
CLIP_OR_DINO_MODEL_PATH="openai/clip-vit-large-patch14-336"
|
||||
SEMANTIC_CONSISTENCY_SAVED_PATH="datasets/panda_70m/videos_clips/meta_consistency_info.jsonl"
|
||||
MIN_SEMANTIC_CONSISTENCY_SCORE=0.8
|
||||
QUALITY_SAVED_PATH="datasets/panda_70m/videos_clips/meta_quality_info_siglip.jsonl"
|
||||
MIN_AESTHETIC_SCORE_SIGLIP=4.0
|
||||
TEXT_SAVED_PATH="datasets/panda_70m/videos_clips/meta_text_info.jsonl"
|
||||
MIN_TEXT_SCORE=0.02
|
||||
MOTION_SAVED_PATH="datasets/panda_70m/videos_clips/meta_motion_info.jsonl"
|
||||
@@ -10,32 +13,55 @@ python -m utils.get_meta_file \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--saved_path $META_FILE_PATH
|
||||
|
||||
# Get the asethetic score (SigLIP) of all videos
|
||||
# Get the semantic consistency score of all video clips.
|
||||
# Adjust the num_workers and batch size parameter based on the machine's computing resources to achieve maximum GPU utilization.
|
||||
accelerate launch compute_semantic_consistency.py \
|
||||
--video_metadata_path $META_FILE_PATH \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--model_path $CLIP_OR_DINO_MODEL_PATH \
|
||||
--frame_sample_method keyframe+first \
|
||||
--batch_size 16 \
|
||||
--num_workers 4 \
|
||||
--saved_freq 10 \
|
||||
--saved_path $SEMANTIC_CONSISTENCY_SAVED_PATH
|
||||
|
||||
# Get the aesthetic score (SigLIP) of all videos filtered by the semantic consistency score.
|
||||
# Adjust the num_workers and batch size parameter based on the machine's computing resources to achieve maximum GPU utilization.
|
||||
accelerate launch compute_video_quality.py \
|
||||
--video_metadata_path $META_FILE_PATH \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--metrics "AestheticScoreSigLIP" \
|
||||
--metrics AestheticScoreSigLIP \
|
||||
--frame_sample_method uniform \
|
||||
--num_sampled_frames 4 \
|
||||
--batch_size 16 \
|
||||
--num_workers 4 \
|
||||
--saved_freq 10 \
|
||||
--saved_path $VIDEO_QUALITY_SAVED_PATH \
|
||||
--batch_size 4
|
||||
--saved_path $QUALITY_SAVED_PATH \
|
||||
--semantic_consistency_score_metadata_path $SEMANTIC_CONSISTENCY_SAVED_PATH \
|
||||
--min_semantic_consistency_score $MIN_SEMANTIC_CONSISTENCY_SCORE
|
||||
|
||||
# Get the text score of all videos filtered by the video quality score.
|
||||
# Get the text score of all videos filtered by the semantic consistency score and video quality score.
|
||||
accelerate launch compute_text_score.py \
|
||||
--video_metadata_path $META_FILE_PATH \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--saved_freq 10 \
|
||||
--saved_path $TEXT_SAVED_PATH \
|
||||
--asethetic_score_siglip_metadata_path $VIDEO_QUALITY_SAVED_PATH \
|
||||
--min_asethetic_score_siglip $MIN_ASETHETIC_SCORE_SIGLIP
|
||||
--semantic_consistency_score_metadata_path $SEMANTIC_CONSISTENCY_SAVED_PATH \
|
||||
--min_semantic_consistency_score $MIN_SEMANTIC_CONSISTENCY_SCORE \
|
||||
--aesthetic_score_siglip_metadata_path $QUALITY_SAVED_PATH \
|
||||
--min_aesthetic_score_siglip $MIN_AESTHETIC_SCORE_SIGLIP
|
||||
|
||||
# Get the motion score of all videos filtered by the video quality score and text score.
|
||||
# Get the motion score of all videos filtered by the semantic consistency score, video quality score and text score.
|
||||
# Adjust the n_jobs parameter based on the actual number of CPU cores in the machine.
|
||||
python compute_motion_score.py \
|
||||
--video_metadata_path $META_FILE_PATH \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--saved_freq 10 \
|
||||
--saved_path $MOTION_SAVED_PATH \
|
||||
--n_jobs 8 \
|
||||
--semantic_consistency_score_metadata_path $SEMANTIC_CONSISTENCY_SAVED_PATH \
|
||||
--min_semantic_consistency_score $MIN_SEMANTIC_CONSISTENCY_SCORE \
|
||||
--aesthetic_score_siglip_metadata_path $QUALITY_SAVED_PATH \
|
||||
--min_aesthetic_score_siglip $MIN_AESTHETIC_SCORE_SIGLIP \
|
||||
--text_score_metadata_path $TEXT_SAVED_PATH \
|
||||
--min_text_score $MIN_TEXT_SCORE
|
||||
--min_text_score $MIN_TEXT_SCORE
|
||||
@@ -1,25 +1,46 @@
|
||||
META_FILE_PATH="datasets/panda_70m/videos_clips/data/meta_file_info.jsonl"
|
||||
VIDEO_FOLDER="datasets/panda_70m/videos_clips/data/"
|
||||
SEMANTIC_CONSISTENCY_SAVED_PATH="datasets/panda_70m/videos_clips/meta_consistency_info.jsonl"
|
||||
MIN_SEMANTIC_CONSISTENCY_SCORE=0.8
|
||||
QUALITY_SAVED_PATH="datasets/panda_70m/videos_clips/meta_quality_info_siglip.jsonl"
|
||||
MIN_AESTHETIC_SCORE_SIGLIP=4.0
|
||||
TEXT_SAVED_PATH="datasets/panda_70m/videos_clips/meta_text_info.jsonl"
|
||||
MIN_TEXT_SCORE=0.02
|
||||
MOTION_SAVED_PATH="datasets/panda_70m/videos_clips/meta_motion_info.jsonl"
|
||||
MIN_MOTION_SCORE=2
|
||||
VIDEO_CAPTION_SAVED_PATH="datasets/panda_70m/meta_caption_info_vila_8b.jsonl"
|
||||
REWRITTEN_VIDEO_CAPTION_SAVED_PATH="datasets/panda_70m/meta_caption_info_vila_8b_rewritten.jsonl"
|
||||
VIDEOCLIPXL_SCORE_SAVED_PATH="datasets/panda_70m/meta_caption_info_vila_8b_rewritten_videoclipxl.jsonl"
|
||||
MAX_MOTION_SCORE=20
|
||||
VIDEO_CAPTION_SAVED_PATH="datasets/panda_70m/meta_caption_info.jsonl"
|
||||
REWRITTEN_VIDEO_CAPTION_SAVED_PATH="datasets/panda_70m/meta_caption_info_rewritten.jsonl"
|
||||
VIDEOCLIPXL_SCORE_SAVED_PATH="datasets/panda_70m/meta_caption_info_rewritten_videoclipxl.jsonl"
|
||||
MIN_VIDEOCLIPXL_SCORE=0.20
|
||||
TRAIN_SAVED_PATH="datasets/panda_70m/train_panda_70m.json"
|
||||
# Manually download Efficient-Large-Model/Llama-3-VILA1.5-8b-AWQ to VILA_MODEL_PATH.
|
||||
# Manually download meta-llama/Meta-Llama-3-8B-Instruct to REWRITE_MODEL_PATH.
|
||||
# Manually download OpenGVLab/InternVL2-40B-AWQ to CAPTION_MODEL_PATH.
|
||||
# You can also download OpenGVLab/InternVL2-2B-AWQ InternVL2-8B-AWQ InternVL2-26B-AWQ or InternVL2-Llama3-76B-AWQ
|
||||
# This a trade-off between recaption quality and speed.
|
||||
CAPTION_MODEL_PATH="OpenGVLab/InternVL2-40B-AWQ"
|
||||
# Manually download meta-llama/Meta-Llama-3.1-70B-Instruct to REWRITE_MODEL_PATH.
|
||||
# You can also download meta-llama/Meta-Llama-3.1-8B-Instruct Meta-Llama-3-8B-Instruct Meta-Llama-3-70B-Instruct.
|
||||
REWRITE_MODEL_PATH="meta-llama/Meta-Llama-3.1-70B-Instruct"
|
||||
|
||||
# Use VILA1.5-AWQ to perform recaptioning.
|
||||
accelerate launch vila_video_recaptioning.py \
|
||||
# Use InternVL2-AWQ to perform recaptioning.
|
||||
# Adjust the num_workers and batch size parameter based on the machine's computing resources to achieve maximum GPU utilization.
|
||||
python3 internvl2_video_recaptioning.py \
|
||||
--video_metadata_path ${META_FILE_PATH} \
|
||||
--video_folder ${VIDEO_FOLDER} \
|
||||
--model_path ${VILA_MODEL_PATH} \
|
||||
--precision "W4A16" \
|
||||
--model_path ${CAPTION_MODEL_PATH} \
|
||||
--saved_path $VIDEO_CAPTION_SAVED_PATH \
|
||||
--saved_freq 1 \
|
||||
--num_workers 4 \
|
||||
--batch_size 128 \
|
||||
--semantic_consistency_score_metadata_path $SEMANTIC_CONSISTENCY_SAVED_PATH \
|
||||
--min_semantic_consistency_score $MIN_SEMANTIC_CONSISTENCY_SCORE \
|
||||
--aesthetic_score_siglip_metadata_path $QUALITY_SAVED_PATH \
|
||||
--min_aesthetic_score_siglip $MIN_AESTHETIC_SCORE_SIGLIP \
|
||||
--text_score_metadata_path $TEXT_SAVED_PATH \
|
||||
--min_text_score $MIN_TEXT_SCORE \
|
||||
--motion_score_metadata_path $MOTION_SAVED_PATH \
|
||||
--min_motion_score $MIN_MOTION_SCORE
|
||||
--min_motion_score $MIN_MOTION_SCORE \
|
||||
--max_motion_score $MAX_MOTION_SCORE
|
||||
|
||||
# Rewrite video captions (optional).
|
||||
python caption_rewrite.py \
|
||||
@@ -32,6 +53,7 @@ python caption_rewrite.py \
|
||||
--saved_freq 1
|
||||
|
||||
# Compute caption-video alignment (optional).
|
||||
# Adjust the num_workers and batch size parameter based on the machine's computing resources to achieve maximum GPU utilization.
|
||||
accelerate launch compute_video_quality.py \
|
||||
--video_metadata_path $REWRITTEN_VIDEO_CAPTION_SAVED_PATH \
|
||||
--caption_column caption \
|
||||
@@ -39,14 +61,15 @@ accelerate launch compute_video_quality.py \
|
||||
--frame_sample_method uniform \
|
||||
--num_sampled_frames 8 \
|
||||
--metrics VideoCLIPXLScore \
|
||||
--batch_size 4 \
|
||||
--num_workers 4 \
|
||||
--batch_size 16 \
|
||||
--saved_path $VIDEOCLIPXL_SCORE_SAVED_PATH \
|
||||
--saved_freq 10
|
||||
|
||||
# Get the final train file.
|
||||
python filter_meta_train.py \
|
||||
--caption_metadata_path $REWRITTEN_VIDEO_CAPTION_SAVED_PATH \
|
||||
--video_folder=$VIDEO_FOLDER \
|
||||
--video_folder $VIDEO_FOLDER \
|
||||
--videoclipxl_score_metadata_path $VIDEOCLIPXL_SCORE_SAVED_PATH \
|
||||
--min_videoclipxl_score $MIN_VIDEOCLIPXL_SCORE \
|
||||
--saved_path=$TRAIN_SAVED_PATH
|
||||
--saved_path $TRAIN_SAVED_PATH
|
||||
@@ -1,5 +1,6 @@
|
||||
import ast
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import pandas as pd
|
||||
|
||||
@@ -7,22 +8,25 @@ from .logger import logger
|
||||
|
||||
|
||||
def filter(
|
||||
video_path_list,
|
||||
basic_metadata_path=None,
|
||||
min_resolution=0,
|
||||
min_duration=-1,
|
||||
max_duration=-1,
|
||||
asethetic_score_metadata_path=None,
|
||||
min_asethetic_score=4,
|
||||
asethetic_score_siglip_metadata_path=None,
|
||||
min_asethetic_score_siglip=4,
|
||||
text_score_metadata_path=None,
|
||||
min_text_score=0.02,
|
||||
motion_score_metadata_path=None,
|
||||
min_motion_score=2,
|
||||
videoclipxl_score_metadata_path=None,
|
||||
min_videoclipxl_score=0.20,
|
||||
video_path_column="video_path",
|
||||
video_path_list: list[str],
|
||||
basic_metadata_path: Optional[str] = None,
|
||||
min_resolution: float = 720*1280,
|
||||
min_duration: float = -1,
|
||||
max_duration: float = -1,
|
||||
aesthetic_score_metadata_path: Optional[str] = None,
|
||||
min_aesthetic_score: float = 4,
|
||||
aesthetic_score_siglip_metadata_path: Optional[str] = None,
|
||||
min_aesthetic_score_siglip: float = 4,
|
||||
text_score_metadata_path: Optional[str] = None,
|
||||
min_text_score: float = 0.02,
|
||||
motion_score_metadata_path: Optional[str] = None,
|
||||
min_motion_score: float = 2,
|
||||
max_motion_score: float = 999999,
|
||||
videoclipxl_score_metadata_path: Optional[str] = None,
|
||||
min_videoclipxl_score: float = 0.20,
|
||||
semantic_consistency_score_metadata_path: Optional[list[str]] = None,
|
||||
min_semantic_consistency_score: float = 0.80,
|
||||
video_path_column: str = "video_path"
|
||||
):
|
||||
video_path_list = [os.path.basename(video_path) for video_path in video_path_list]
|
||||
|
||||
@@ -31,7 +35,7 @@ def filter(
|
||||
basic_df = pd.read_csv(basic_metadata_path)
|
||||
elif basic_metadata_path.endswith(".jsonl"):
|
||||
basic_df = pd.read_json(basic_metadata_path, lines=True)
|
||||
|
||||
|
||||
basic_df["resolution"] = basic_df["frame_size"].apply(lambda x: x[0] * x[1])
|
||||
filtered_basic_df = basic_df[basic_df["resolution"] < min_resolution]
|
||||
filtered_video_path_list = filtered_basic_df[video_path_column].tolist()
|
||||
@@ -65,50 +69,50 @@ def filter(
|
||||
f"with duration greater than {max_duration}."
|
||||
)
|
||||
|
||||
if asethetic_score_metadata_path is not None:
|
||||
if asethetic_score_metadata_path.endswith(".csv"):
|
||||
asethetic_score_df = pd.read_csv(asethetic_score_metadata_path)
|
||||
elif asethetic_score_metadata_path.endswith(".jsonl"):
|
||||
asethetic_score_df = pd.read_json(asethetic_score_metadata_path, lines=True)
|
||||
if aesthetic_score_metadata_path is not None:
|
||||
if aesthetic_score_metadata_path.endswith(".csv"):
|
||||
aesthetic_score_df = pd.read_csv(aesthetic_score_metadata_path)
|
||||
elif aesthetic_score_metadata_path.endswith(".jsonl"):
|
||||
aesthetic_score_df = pd.read_json(aesthetic_score_metadata_path, lines=True)
|
||||
|
||||
# In pandas, csv will save lists as strings, whereas jsonl will not.
|
||||
asethetic_score_df["aesthetic_score"] = asethetic_score_df["aesthetic_score"].apply(
|
||||
aesthetic_score_df["aesthetic_score"] = aesthetic_score_df["aesthetic_score"].apply(
|
||||
lambda x: ast.literal_eval(x) if isinstance(x, str) else x
|
||||
)
|
||||
asethetic_score_df["aesthetic_score_mean"] = asethetic_score_df["aesthetic_score"].apply(lambda x: sum(x) / len(x))
|
||||
filtered_asethetic_score_df = asethetic_score_df[asethetic_score_df["aesthetic_score_mean"] < min_asethetic_score]
|
||||
filtered_video_path_list = filtered_asethetic_score_df[video_path_column].tolist()
|
||||
aesthetic_score_df["aesthetic_score_mean"] = aesthetic_score_df["aesthetic_score"].apply(lambda x: sum(x) / len(x))
|
||||
filtered_aesthetic_score_df = aesthetic_score_df[aesthetic_score_df["aesthetic_score_mean"] < min_aesthetic_score]
|
||||
filtered_video_path_list = filtered_aesthetic_score_df[video_path_column].tolist()
|
||||
filtered_video_path_list = [os.path.basename(video_path) for video_path in filtered_video_path_list]
|
||||
|
||||
video_path_list = list(set(video_path_list).difference(set(filtered_video_path_list)))
|
||||
logger.info(
|
||||
f"Load {asethetic_score_metadata_path} ({len(asethetic_score_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with aesthetic score less than {min_asethetic_score}."
|
||||
f"Load {aesthetic_score_metadata_path} ({len(aesthetic_score_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with aesthetic score less than {min_aesthetic_score}."
|
||||
)
|
||||
|
||||
if asethetic_score_siglip_metadata_path is not None:
|
||||
if asethetic_score_siglip_metadata_path.endswith(".csv"):
|
||||
asethetic_score_siglip_df = pd.read_csv(asethetic_score_siglip_metadata_path)
|
||||
elif asethetic_score_siglip_metadata_path.endswith(".jsonl"):
|
||||
asethetic_score_siglip_df = pd.read_json(asethetic_score_siglip_metadata_path, lines=True)
|
||||
|
||||
if aesthetic_score_siglip_metadata_path is not None:
|
||||
if aesthetic_score_siglip_metadata_path.endswith(".csv"):
|
||||
aesthetic_score_siglip_df = pd.read_csv(aesthetic_score_siglip_metadata_path)
|
||||
elif aesthetic_score_siglip_metadata_path.endswith(".jsonl"):
|
||||
aesthetic_score_siglip_df = pd.read_json(aesthetic_score_siglip_metadata_path, lines=True)
|
||||
|
||||
# In pandas, csv will save lists as strings, whereas jsonl will not.
|
||||
asethetic_score_siglip_df["aesthetic_score_siglip"] = asethetic_score_siglip_df["aesthetic_score_siglip"].apply(
|
||||
aesthetic_score_siglip_df["aesthetic_score_siglip"] = aesthetic_score_siglip_df["aesthetic_score_siglip"].apply(
|
||||
lambda x: ast.literal_eval(x) if isinstance(x, str) else x
|
||||
)
|
||||
asethetic_score_siglip_df["aesthetic_score_siglip_mean"] = asethetic_score_siglip_df["aesthetic_score_siglip"].apply(
|
||||
aesthetic_score_siglip_df["aesthetic_score_siglip_mean"] = aesthetic_score_siglip_df["aesthetic_score_siglip"].apply(
|
||||
lambda x: sum(x) / len(x)
|
||||
)
|
||||
filtered_asethetic_score_siglip_df = asethetic_score_siglip_df[
|
||||
asethetic_score_siglip_df["aesthetic_score_siglip_mean"] < min_asethetic_score_siglip
|
||||
filtered_aesthetic_score_siglip_df = aesthetic_score_siglip_df[
|
||||
aesthetic_score_siglip_df["aesthetic_score_siglip_mean"] < min_aesthetic_score_siglip
|
||||
]
|
||||
filtered_video_path_list = filtered_asethetic_score_siglip_df[video_path_column].tolist()
|
||||
filtered_video_path_list = filtered_aesthetic_score_siglip_df[video_path_column].tolist()
|
||||
filtered_video_path_list = [os.path.basename(video_path) for video_path in filtered_video_path_list]
|
||||
|
||||
video_path_list = list(set(video_path_list).difference(set(filtered_video_path_list)))
|
||||
logger.info(
|
||||
f"Load {asethetic_score_siglip_metadata_path} ({len(asethetic_score_siglip_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with aesthetic score (SigLIP) less than {min_asethetic_score_siglip}."
|
||||
f"Load {aesthetic_score_siglip_metadata_path} ({len(aesthetic_score_siglip_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with aesthetic score (SigLIP) less than {min_aesthetic_score_siglip}."
|
||||
)
|
||||
|
||||
if text_score_metadata_path is not None:
|
||||
@@ -132,7 +136,7 @@ def filter(
|
||||
motion_score_df = pd.read_csv(motion_score_metadata_path)
|
||||
elif motion_score_metadata_path.endswith(".jsonl"):
|
||||
motion_score_df = pd.read_json(motion_score_metadata_path, lines=True)
|
||||
|
||||
|
||||
filtered_motion_score_df = motion_score_df[motion_score_df["motion_score"] < min_motion_score]
|
||||
filtered_video_path_list = filtered_motion_score_df[video_path_column].tolist()
|
||||
filtered_video_path_list = [os.path.basename(video_path) for video_path in filtered_video_path_list]
|
||||
@@ -142,6 +146,16 @@ def filter(
|
||||
f"Load {motion_score_metadata_path} ({len(motion_score_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with motion score smaller than {min_motion_score}."
|
||||
)
|
||||
|
||||
filtered_motion_score_df = motion_score_df[motion_score_df["motion_score"] > max_motion_score]
|
||||
filtered_video_path_list = filtered_motion_score_df[video_path_column].tolist()
|
||||
filtered_video_path_list = [os.path.basename(video_path) for video_path in filtered_video_path_list]
|
||||
|
||||
video_path_list = list(set(video_path_list).difference(set(filtered_video_path_list)))
|
||||
logger.info(
|
||||
f"Load {motion_score_metadata_path} ({len(motion_score_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with motion score greater than {min_motion_score}."
|
||||
)
|
||||
|
||||
if videoclipxl_score_metadata_path is not None:
|
||||
if videoclipxl_score_metadata_path.endswith(".csv"):
|
||||
@@ -158,5 +172,23 @@ def filter(
|
||||
f"Load {videoclipxl_score_metadata_path} ({len(videoclipxl_score_df)}) and "
|
||||
f"filter {len(filtered_video_path_list)} videos with mixclip score smaller than {min_videoclipxl_score}."
|
||||
)
|
||||
|
||||
if semantic_consistency_score_metadata_path is not None:
|
||||
for f in semantic_consistency_score_metadata_path:
|
||||
if f.endswith(".csv"):
|
||||
semantic_consistency_score_df = pd.read_csv(f)
|
||||
elif f.endswith(".jsonl"):
|
||||
semantic_consistency_score_df = pd.read_json(f, lines=True)
|
||||
filtered_semantic_consistency_score_df = semantic_consistency_score_df[
|
||||
semantic_consistency_score_df["similarity_cross_frame"].apply(lambda x: min(x) < min_semantic_consistency_score)
|
||||
]
|
||||
filtered_video_path_list = filtered_semantic_consistency_score_df[video_path_column].tolist()
|
||||
filtered_video_path_list = [os.path.basename(video_path) for video_path in filtered_video_path_list]
|
||||
|
||||
video_path_list = list(set(video_path_list).difference(set(filtered_video_path_list)))
|
||||
logger.info(
|
||||
f"Load {f} ({len(semantic_consistency_score_df)}) and filter {len(filtered_video_path_list)} videos "
|
||||
f"with the minimum semantic consistency score smaller than {min_semantic_consistency_score}."
|
||||
)
|
||||
|
||||
return video_path_list
|
||||
|
||||
@@ -9,7 +9,8 @@ from torch.utils.data import DataLoader, Dataset
|
||||
from .logger import logger
|
||||
from .video_utils import extract_frames
|
||||
|
||||
ALL_VIDEO_EXT = set(["mp4", "webm", "mkv", "avi", "flv", "mov"])
|
||||
|
||||
ALL_VIDEO_EXT = set([".mp4", ".webm", ".mkv", ".avi", ".flv", ".mov", ".ts"])
|
||||
VIDEO_READER_TIMEOUT = 300
|
||||
|
||||
|
||||
@@ -29,7 +30,7 @@ class VideoDataset(Dataset):
|
||||
text_column: Optional[str] = None,
|
||||
sample_method: str = "mid",
|
||||
num_sampled_frames: int = 1,
|
||||
num_sample_stride: Optional[int] = None
|
||||
sample_stride: Optional[int] = None
|
||||
):
|
||||
length = len(dataset_inputs[list(dataset_inputs.keys())[0]])
|
||||
if not all(len(v) == length for v in dataset_inputs.values()):
|
||||
@@ -45,7 +46,7 @@ class VideoDataset(Dataset):
|
||||
|
||||
self.sample_method = sample_method
|
||||
self.num_sampled_frames = num_sampled_frames
|
||||
self.num_sample_stride = num_sample_stride
|
||||
self.sample_stride = sample_stride
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_path = self.video_path_list[index]
|
||||
@@ -60,7 +61,7 @@ class VideoDataset(Dataset):
|
||||
else:
|
||||
# It is a trick to deal with decord hanging when reading some abnormal videos.
|
||||
try:
|
||||
sample_args = (video_path, self.sample_method, self.num_sampled_frames, self.num_sample_stride)
|
||||
sample_args = (video_path, self.sample_method, self.num_sampled_frames, self.sample_stride)
|
||||
sampled_frame_idx_list, sampled_frame_list = func_timeout(
|
||||
VIDEO_READER_TIMEOUT, extract_frames, args=sample_args
|
||||
)
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import gc
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
from contextlib import contextmanager
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
@@ -8,6 +10,11 @@ from decord import VideoReader
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ALL_FRAME_SAMPLE_METHODS = [
|
||||
"mid", "uniform", "random", "stride", "first", "last", "keyframe", "keyframe+first", "keyframe+last"
|
||||
]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def video_reader(*args, **kwargs):
|
||||
"""A context manager to solve the memory leak of decord.
|
||||
@@ -20,13 +27,44 @@ def video_reader(*args, **kwargs):
|
||||
gc.collect()
|
||||
|
||||
|
||||
def get_keyframe_index(video_path):
|
||||
"""Extract the frame index list of I-frames. In general, the first frame in a video should be the I-frame.
|
||||
The extracted frame index is more accurate than the pts_time * avg_fps.
|
||||
"""
|
||||
assert shutil.which("ffprobe") is not None, f"Please install ffprobe and make sure it is in the system path."
|
||||
|
||||
command = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "frame=pict_type",
|
||||
"-of", "csv=p=0",
|
||||
video_path
|
||||
]
|
||||
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
|
||||
|
||||
keyframe_index_list = []
|
||||
for index, line in enumerate(result.stdout.split("\n")):
|
||||
line = line.strip(",")
|
||||
pict_type = line.strip()
|
||||
if pict_type == "I":
|
||||
keyframe_index_list.append(index)
|
||||
|
||||
return keyframe_index_list
|
||||
|
||||
def extract_frames(
|
||||
video_path: str,
|
||||
sample_method: str = "mid",
|
||||
num_sampled_frames: int = -1,
|
||||
sample_stride: int = -1,
|
||||
num_sampled_frames: int = 1,
|
||||
sample_stride: Optional[int] = None,
|
||||
**kwargs
|
||||
) -> Optional[Tuple[List[int], List[Image.Image]]]:
|
||||
if num_sampled_frames < 1:
|
||||
raise ValueError(f"The num_sampled_frames must be greater than 1.")
|
||||
if sample_stride is not None and sample_stride < 1:
|
||||
raise ValueError(f"The sample_stride must be greater than 1.")
|
||||
if sample_stride is not None and sample_method not in ["random", "stride"]:
|
||||
raise ValueError(f"The sample_method must be random or stride when sample_stride is specified.")
|
||||
with video_reader(video_path, num_threads=2, **kwargs) as vr:
|
||||
if sample_method == "mid":
|
||||
sampled_frame_idx_list = [len(vr) // 2]
|
||||
@@ -36,8 +74,24 @@ def extract_frames(
|
||||
clip_length = min(len(vr), (num_sampled_frames - 1) * sample_stride + 1)
|
||||
start_idx = random.randint(0, len(vr) - clip_length)
|
||||
sampled_frame_idx_list = np.linspace(start_idx, start_idx + clip_length - 1, num_sampled_frames, dtype=int)
|
||||
elif sample_method == "stride":
|
||||
sampled_frame_idx_list = np.arange(0, len(vr), sample_stride)
|
||||
elif sample_method == "first":
|
||||
sampled_frame_idx_list = [0]
|
||||
elif sample_method == "last":
|
||||
sampled_frame_idx_list = [len(vr) - 1]
|
||||
elif sample_method == "keyframe":
|
||||
sampled_frame_idx_list = get_keyframe_index(video_path)
|
||||
elif sample_method == "keyframe+first":
|
||||
sampled_frame_idx_list = get_keyframe_index(video_path)
|
||||
if len(sampled_frame_idx_list) == 1 or sampled_frame_idx_list[1] > 1 * vr.get_avg_fps():
|
||||
sampled_frame_idx_list.insert(1, int(1 * vr.get_avg_fps()))
|
||||
elif sample_method == "keyframe+last":
|
||||
sampled_frame_idx_list = get_keyframe_index(video_path)
|
||||
if sampled_frame_idx_list[-1] != (len(vr) - 1):
|
||||
sampled_frame_idx_list.append(len(vr) - 1)
|
||||
else:
|
||||
raise ValueError(f"The sample_method {sample_method} must be mid, uniform or random.")
|
||||
raise ValueError(f"The sample_method must be within {ALL_FRAME_SAMPLE_METHODS}.")
|
||||
sampled_frame_list = vr.get_batch(sampled_frame_idx_list).asnumpy()
|
||||
sampled_frame_list = [Image.fromarray(frame) for frame in sampled_frame_list]
|
||||
|
||||
|
||||
@@ -1,343 +0,0 @@
|
||||
# Modified from https://github.com/mit-han-lab/llm-awq/blob/main/tinychat/vlm_demo_new.py.
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import tinychat.utils.constants
|
||||
import torch
|
||||
from accelerate import PartialState, load_checkpoint_and_dispatch
|
||||
from accelerate.utils import gather_object
|
||||
from decord import VideoReader
|
||||
from natsort import natsorted
|
||||
from PIL import Image
|
||||
# from tinychat.models.llava_llama import LlavaLlamaForCausalLM
|
||||
from tinychat.models.vila_llama import VilaLlamaForCausalLM
|
||||
from tinychat.stream_generators.llava_stream_gen import LlavaStreamGenerator
|
||||
from tinychat.utils.conversation_utils import gen_params
|
||||
from tinychat.utils.llava_image_processing import process_images
|
||||
from tinychat.utils.prompt_templates import (get_image_token, get_prompter,
|
||||
get_stop_token_ids)
|
||||
from tinychat.utils.tune import device_warmup, tune_llava_patch_embedding
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoConfig, AutoTokenizer
|
||||
from utils.filter import filter
|
||||
from utils.logger import logger
|
||||
|
||||
gen_params.seed = 1
|
||||
gen_params.temp = 1.0
|
||||
gen_params.top_p = 1.0
|
||||
|
||||
|
||||
def extract_uniform_frames(video_path: str, num_sampled_frames: int = 8):
|
||||
vr = VideoReader(video_path)
|
||||
sampled_frame_idx_list = np.linspace(0, len(vr), num_sampled_frames, endpoint=False, dtype=int)
|
||||
sampled_frame_list = []
|
||||
for idx in sampled_frame_idx_list:
|
||||
sampled_frame = Image.fromarray(vr[idx].asnumpy())
|
||||
sampled_frame_list.append(sampled_frame)
|
||||
|
||||
return sampled_frame_list
|
||||
|
||||
|
||||
def stream_output(output_stream):
|
||||
for outputs in output_stream:
|
||||
output_text = outputs["text"]
|
||||
output_text = output_text.strip().split(" ")
|
||||
# print(f"output_text: {output_text}.")
|
||||
return " ".join(output_text)
|
||||
|
||||
|
||||
def skip(*args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Recaption videos with VILA1.5.")
|
||||
parser.add_argument(
|
||||
"--video_metadata_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="The path to the video dataset metadata (csv/jsonl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_path_column",
|
||||
type=str,
|
||||
default="video_path",
|
||||
help="The column contains the video path (an absolute path or a relative path w.r.t the video_folder).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--caption_column",
|
||||
type=str,
|
||||
default="caption",
|
||||
help="The column contains the caption.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video_folder", type=str, default="", help="The video folder."
|
||||
)
|
||||
parser.add_argument("--input_prompt", type=str, default="<video>\\n Elaborate on the visual and narrative elements of the video in detail.")
|
||||
parser.add_argument(
|
||||
"--model_type", type=str, default="LLaMa", help="type of the model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_path", type=str, default="Efficient-Large-Model/Llama-3-VILA1.5-8b-AWQ"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quant_path",
|
||||
type=str,
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--precision", type=str, default="W4A16", help="compute precision"
|
||||
)
|
||||
parser.add_argument("--num_sampled_frames", type=int, default=8)
|
||||
parser.add_argument(
|
||||
"--saved_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The save path to the output results (csv/jsonl).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--saved_freq",
|
||||
type=int,
|
||||
default=100,
|
||||
help="The frequency to save the output results.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--basic_metadata_path", type=str, default=None, help="The path to the basic metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_resolution", type=float, default=0, help="The resolution threshold.")
|
||||
parser.add_argument("--min_duration", type=float, default=-1, help="The minimum duration.")
|
||||
parser.add_argument("--max_duration", type=float, default=-1, help="The maximum duration.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score", type=float, default=4.0, help="The asethetic score threshold.")
|
||||
parser.add_argument(
|
||||
"--asethetic_score_siglip_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_asethetic_score_siglip", type=float, default=4.0, help="The asethetic score (SigLIP) threshold.")
|
||||
parser.add_argument(
|
||||
"--text_score_metadata_path", type=str, default=None, help="The path to the video text score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_text_score", type=float, default=0.02, help="The text threshold.")
|
||||
parser.add_argument(
|
||||
"--motion_score_metadata_path", type=str, default=None, help="The path to the video motion score metadata (csv/jsonl)."
|
||||
)
|
||||
parser.add_argument("--min_motion_score", type=float, default=2, help="The motion threshold.")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main(args):
|
||||
if args.video_metadata_path.endswith(".csv"):
|
||||
video_metadata_df = pd.read_csv(args.video_metadata_path)
|
||||
elif args.video_metadata_path.endswith(".jsonl"):
|
||||
video_metadata_df = pd.read_json(args.video_metadata_path, lines=True)
|
||||
else:
|
||||
raise ValueError("The video_metadata_path must end with .csv or .jsonl.")
|
||||
video_path_list = video_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = [os.path.basename(video_path) for video_path in video_path_list]
|
||||
|
||||
if not (args.saved_path.endswith(".csv") or args.saved_path.endswith(".jsonl")):
|
||||
raise ValueError("The saved_path must end with .csv or .jsonl.")
|
||||
|
||||
if os.path.exists(args.saved_path):
|
||||
if args.saved_path.endswith(".csv"):
|
||||
saved_metadata_df = pd.read_csv(args.saved_path)
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
saved_metadata_df = pd.read_json(args.saved_path, lines=True)
|
||||
saved_video_path_list = saved_metadata_df[args.video_path_column].tolist()
|
||||
video_path_list = list(set(video_path_list).difference(set(saved_video_path_list)))
|
||||
logger.info(
|
||||
f"Resume from {args.saved_path}: {len(saved_video_path_list)} processed and {len(video_path_list)} to be processed."
|
||||
)
|
||||
|
||||
video_path_list = filter(
|
||||
video_path_list,
|
||||
basic_metadata_path=args.basic_metadata_path,
|
||||
min_resolution=args.min_resolution,
|
||||
min_duration=args.min_duration,
|
||||
max_duration=args.max_duration,
|
||||
asethetic_score_metadata_path=args.asethetic_score_metadata_path,
|
||||
min_asethetic_score=args.min_asethetic_score,
|
||||
asethetic_score_siglip_metadata_path=args.asethetic_score_siglip_metadata_path,
|
||||
min_asethetic_score_siglip=args.min_asethetic_score_siglip,
|
||||
text_score_metadata_path=args.text_score_metadata_path,
|
||||
min_text_score=args.min_text_score,
|
||||
motion_score_metadata_path=args.motion_score_metadata_path,
|
||||
min_motion_score=args.min_motion_score,
|
||||
)
|
||||
video_path_list = [os.path.join(args.video_folder, video_path) for video_path in video_path_list]
|
||||
# Sorting to guarantee the same result for each process.
|
||||
video_path_list = natsorted(video_path_list)
|
||||
|
||||
state = PartialState()
|
||||
|
||||
# Accelerate model initialization
|
||||
setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
|
||||
setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
|
||||
torch.nn.init.kaiming_uniform_ = skip
|
||||
torch.nn.init.kaiming_normal_ = skip
|
||||
torch.nn.init.uniform_ = skip
|
||||
torch.nn.init.normal_ = skip
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(os.path.join(args.model_path, "llm"), use_fast=False)
|
||||
tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
|
||||
tokenizer.convert_tokens_to_ids(
|
||||
[tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
|
||||
)[0]
|
||||
)
|
||||
config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True)
|
||||
model = VilaLlamaForCausalLM(config).half()
|
||||
tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN_IDX = (
|
||||
tokenizer.convert_tokens_to_ids(
|
||||
[tinychat.utils.constants.LLAVA_DEFAULT_IMAGE_PATCH_TOKEN]
|
||||
)[0]
|
||||
)
|
||||
vision_tower = model.get_vision_tower()
|
||||
# if not vision_tower.is_loaded:
|
||||
# vision_tower.load_model()
|
||||
image_processor = vision_tower.image_processor
|
||||
# vision_tower = vision_tower.half()
|
||||
|
||||
if args.precision == "W16A16":
|
||||
pbar = tqdm(range(1))
|
||||
pbar.set_description("Loading checkpoint shards")
|
||||
for i in pbar:
|
||||
model.llm = load_checkpoint_and_dispatch(
|
||||
model.llm,
|
||||
os.path.join(args.model_path, "llm"),
|
||||
no_split_module_classes=[
|
||||
"OPTDecoderLayer",
|
||||
"LlamaDecoderLayer",
|
||||
"BloomBlock",
|
||||
"MPTBlock",
|
||||
"DecoderLayer",
|
||||
"CLIPEncoderLayer",
|
||||
],
|
||||
).to(state.device)
|
||||
model = model.to(state.device)
|
||||
|
||||
elif args.precision == "W4A16":
|
||||
from tinychat.utils.load_quant import load_awq_model
|
||||
|
||||
# Auto load quant_path from the 3b/8b/13b/40b model.
|
||||
if args.quant_path is None:
|
||||
if "VILA1.5-3b-s2-AWQ" in args.model_path:
|
||||
args.quant_path = os.path.join(args.model_path, "llm/vila-1.5-3b-s2-w4-g128-awq-v2.pt")
|
||||
elif "VILA1.5-3b-AWQ" in args.model_path:
|
||||
args.quant_path = os.path.join(args.model_path, "llm/vila-1.5-3b-w4-g128-awq-v2.pt")
|
||||
elif "Llama-3-VILA1.5-8b-AWQ" in args.model_path:
|
||||
args.quant_path = os.path.join(args.model_path, "llm/llama-3-vila1.5-8b-w4-g128-awq-v2.pt")
|
||||
elif "VILA1.5-13b-AWQ" in args.model_path:
|
||||
args.quant_path = os.path.join(args.model_path, "llm/vila-1.5-13b-w4-g128-awq-v2.pt")
|
||||
elif "VILA1.5-40b-AWQ" in args.model_path:
|
||||
args.quant_path = os.path.join(args.model_path, "llm/vila-1.5-40b-w4-g128-awq-v2.pt")
|
||||
model.llm = load_awq_model(model.llm, args.quant_path, 4, 128, state.device)
|
||||
from tinychat.modules import (make_fused_mlp, make_fused_vision_attn,
|
||||
make_quant_attn, make_quant_norm)
|
||||
|
||||
make_quant_attn(model.llm, state.device)
|
||||
make_quant_norm(model.llm)
|
||||
# make_fused_mlp(model)
|
||||
# make_fused_vision_attn(model,state.device)
|
||||
model = model.to(state.device)
|
||||
|
||||
else:
|
||||
raise NotImplementedError(f"Precision {args.precision} is not supported.")
|
||||
|
||||
device_warmup(state.device)
|
||||
tune_llava_patch_embedding(vision_tower, device=state.device)
|
||||
|
||||
stream_generator = LlavaStreamGenerator
|
||||
|
||||
model_prompter = get_prompter(
|
||||
args.model_type, args.model_path, False, False
|
||||
)
|
||||
stop_token_ids = get_stop_token_ids(args.model_type, args.model_path)
|
||||
|
||||
model.eval()
|
||||
|
||||
index = len(video_path_list) - len(video_path_list) % state.num_processes
|
||||
# Avoid the NCCL timeout in the final gather operation.
|
||||
logger.info(f"Drop {len(video_path_list) % state.num_processes} videos to ensure each process handles the same number of videos.")
|
||||
video_path_list = video_path_list[:index]
|
||||
logger.info(f"{len(video_path_list)} videos are to be processed.")
|
||||
|
||||
result_dict = {args.video_path_column: [], args.caption_column: []}
|
||||
with state.split_between_processes(video_path_list) as splitted_video_path_list:
|
||||
# TODO: Use VideoDataset.
|
||||
for i, video_path in enumerate(tqdm(splitted_video_path_list)):
|
||||
try:
|
||||
image_list = extract_uniform_frames(video_path, args.num_sampled_frames)
|
||||
image_num = len(image_list)
|
||||
# Similar operation in model_worker.py
|
||||
image_tensor = process_images(image_list, image_processor, model.config)
|
||||
if type(image_tensor) is list:
|
||||
image_tensor = [
|
||||
image.to(state.device, dtype=torch.float16) for image in image_tensor
|
||||
]
|
||||
else:
|
||||
image_tensor = image_tensor.to(state.device, dtype=torch.float16)
|
||||
|
||||
input_prompt = args.input_prompt
|
||||
# Insert image here
|
||||
image_token = get_image_token(model, args.model_path)
|
||||
image_token_holder = tinychat.utils.constants.LLAVA_DEFAULT_IM_TOKEN_PLACE_HOLDER
|
||||
im_token_count = input_prompt.count(image_token_holder)
|
||||
if im_token_count == 0:
|
||||
model_prompter.insert_prompt(image_token * image_num + input_prompt)
|
||||
else:
|
||||
assert im_token_count == image_num
|
||||
input_prompt = input_prompt.replace(image_token_holder, image_token)
|
||||
model_prompter.insert_prompt(input_prompt)
|
||||
output_stream = stream_generator(
|
||||
model,
|
||||
tokenizer,
|
||||
model_prompter.model_input,
|
||||
gen_params,
|
||||
device=state.device,
|
||||
stop_token_ids=stop_token_ids,
|
||||
image_tensor=image_tensor,
|
||||
)
|
||||
outputs = stream_output(output_stream)
|
||||
if len(outputs) != 0:
|
||||
result_dict[args.video_path_column].append(Path(video_path).name)
|
||||
result_dict[args.caption_column].append(outputs)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"VILA with {video_path} failed. Error is {e}.")
|
||||
|
||||
if i != 0 and i % args.saved_freq == 0:
|
||||
state.wait_for_everyone()
|
||||
gathered_result_dict = {k: gather_object(v) for k, v in result_dict.items()}
|
||||
if state.is_main_process and len(gathered_result_dict[args.video_path_column]) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_dict)
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
for k in result_dict.keys():
|
||||
result_dict[k] = []
|
||||
|
||||
state.wait_for_everyone()
|
||||
gathered_result_dict = {k: gather_object(v) for k, v in result_dict.items()}
|
||||
if state.is_main_process and len(gathered_result_dict[args.video_path_column]) != 0:
|
||||
result_df = pd.DataFrame(gathered_result_dict)
|
||||
if args.saved_path.endswith(".csv"):
|
||||
header = False if os.path.exists(args.saved_path) else True
|
||||
result_df.to_csv(args.saved_path, header=header, index=False, mode="a")
|
||||
elif args.saved_path.endswith(".jsonl"):
|
||||
result_df.to_json(args.saved_path, orient="records", lines=True, mode="a", force_ascii=False)
|
||||
logger.info(f"Save result to {args.saved_path}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
Reference in New Issue
Block a user