134 lines
5.8 KiB
Python
134 lines
5.8 KiB
Python
import argparse
|
|
import os
|
|
import re
|
|
from tqdm import tqdm
|
|
|
|
import pandas as pd
|
|
from vllm import LLM, SamplingParams
|
|
|
|
from utils.logger import logger
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Recaption the video frame.")
|
|
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(
|
|
"--caption_column",
|
|
type=str,
|
|
default="sampled_frame_caption",
|
|
help="The column contains the sampled_frame_caption.",
|
|
)
|
|
parser.add_argument(
|
|
"--remove_quotes",
|
|
action="store_true",
|
|
help="Whether to remove quotes from caption.",
|
|
)
|
|
parser.add_argument(
|
|
"--batch_size",
|
|
type=int,
|
|
default=10,
|
|
required=False,
|
|
help="The batch size for the video caption.",
|
|
)
|
|
parser.add_argument(
|
|
"--summary_model_name",
|
|
type=str,
|
|
default="mistralai/Mistral-7B-Instruct-v0.2",
|
|
)
|
|
parser.add_argument(
|
|
"--summary_prompt",
|
|
type=str,
|
|
default=(
|
|
"You are a helpful video description generator. I'll give you a description of the middle frame of the video clip, "
|
|
"which you need to summarize it into a description of the video clip."
|
|
"Please provide your video description following these requirements: "
|
|
"1. Describe the basic and necessary information of the video in the third person, be as concise as possible. "
|
|
"2. Output the video description directly. Begin with 'In this video'. "
|
|
"3. Limit the video description within 100 words. "
|
|
"Here is the mid-frame description: "
|
|
),
|
|
)
|
|
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.")
|
|
|
|
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()
|
|
sampled_frame_caption_list = video_metadata_df[args.caption_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) - set(saved_video_path_list))
|
|
video_metadata_df.set_index(args.video_path_column, inplace=True)
|
|
video_metadata_df = video_metadata_df.loc[video_path_list]
|
|
sampled_frame_caption_list = video_metadata_df[args.caption_column].tolist()
|
|
logger.info(f"Resume from {args.saved_path}: {len(saved_video_path_list)} processed and {len(video_path_list)} to be processed.")
|
|
|
|
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=256)
|
|
summary_model = LLM(model=args.summary_model_name, trust_remote_code=True)
|
|
|
|
result_dict = {"video_path": [], "summary_model": [], "summary_caption": []}
|
|
|
|
for i in tqdm(range(0, len(sampled_frame_caption_list), args.batch_size)):
|
|
batch_video_path = video_path_list[i: i + args.batch_size]
|
|
batch_caption = sampled_frame_caption_list[i : i + args.batch_size]
|
|
batch_prompt = []
|
|
for caption in batch_caption:
|
|
if args.remove_quotes:
|
|
caption = re.sub(r'(["\']).*?\1', "", caption)
|
|
batch_prompt.append("user:" + args.summary_prompt + str(caption) + "\n assistant:")
|
|
batch_output = summary_model.generate(batch_prompt, sampling_params)
|
|
|
|
result_dict["video_path"].extend(batch_video_path)
|
|
result_dict["summary_model"].extend([args.summary_model_name] * len(batch_caption))
|
|
result_dict["summary_caption"].extend([output.outputs[0].text.rstrip() for output in batch_output])
|
|
|
|
# Save the metadata every args.saved_freq.
|
|
if i != 0 and ((i // args.batch_size) % args.saved_freq) == 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")
|
|
logger.info(f"Save result to {args.saved_path}.")
|
|
|
|
result_dict = {"video_path": [], "summary_model": [], "summary_caption": []}
|
|
|
|
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")
|
|
logger.info(f"Save the final result to {args.saved_path}.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |