Spaces:
Running
Running
| import ast | |
| import argparse | |
| import os | |
| from pathlib import Path | |
| import easyocr | |
| import numpy as np | |
| import pandas as pd | |
| from accelerate import PartialState | |
| from accelerate.utils import gather_object | |
| from natsort import natsorted | |
| from tqdm import tqdm | |
| from torchvision.datasets.utils import download_url | |
| from utils.logger import logger | |
| from utils.video_utils import extract_frames, get_video_path_list | |
| def init_ocr_reader(root: str = "~/.cache/easyocr", device: str = "gpu"): | |
| root = os.path.expanduser(root) | |
| if not os.path.exists(root): | |
| os.makedirs(root) | |
| download_url( | |
| "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/easyanimate/video_caption/easyocr/craft_mlt_25k.pth", | |
| root, | |
| filename="craft_mlt_25k.pth", | |
| md5="2f8227d2def4037cdb3b34389dcf9ec1", | |
| ) | |
| ocr_reader = easyocr.Reader( | |
| lang_list=["en", "ch_sim"], | |
| gpu=device, | |
| recognizer=False, | |
| verbose=False, | |
| model_storage_directory=root, | |
| ) | |
| return ocr_reader | |
| def triangle_area(p1, p2, p3): | |
| """Compute the triangle area according to its coordinates. | |
| """ | |
| x1, y1 = p1 | |
| x2, y2 = p2 | |
| x3, y3 = p3 | |
| tri_area = 0.5 * np.abs(x1 * y2 + x2 * y3 + x3 * y1 - x2 * y1 - x3 * y2 - x1 * y3) | |
| return tri_area | |
| def compute_text_score(video_path, ocr_reader): | |
| _, images = extract_frames(video_path, sample_method="mid") | |
| images = [np.array(image) for image in images] | |
| frame_ocr_area_ratios = [] | |
| for image in images: | |
| # horizontal detected results and free-form detected | |
| horizontal_list, free_list = ocr_reader.detect(np.asarray(image)) | |
| width, height = image.shape[0], image.shape[1] | |
| total_area = width * height | |
| # rectangles | |
| rect_area = 0 | |
| for xmin, xmax, ymin, ymax in horizontal_list[0]: | |
| if xmax < xmin or ymax < ymin: | |
| continue | |
| rect_area += (xmax - xmin) * (ymax - ymin) | |
| # free-form | |
| quad_area = 0 | |
| try: | |
| for points in free_list[0]: | |
| triangle1 = points[:3] | |
| quad_area += triangle_area(*triangle1) | |
| triangle2 = points[3:] + [points[0]] | |
| quad_area += triangle_area(*triangle2) | |
| except: | |
| quad_area = 0 | |
| text_area = rect_area + quad_area | |
| frame_ocr_area_ratios.append(text_area / total_area) | |
| video_meta_info = { | |
| "video_path": Path(video_path).name, | |
| "text_score": round(np.mean(frame_ocr_area_ratios), 5), | |
| } | |
| return video_meta_info | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Compute the text score of the middle frame in the videos.") | |
| parser.add_argument("--video_folder", type=str, default="", help="The video folder.") | |
| 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("--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( | |
| "--asethetic_score_metadata_path", type=str, default=None, help="The path to the video quality metadata (csv/jsonl)." | |
| ) | |
| parser.add_argument("--asethetic_score_threshold", type=float, default=4.0, help="The asethetic score threshold.") | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = parse_args() | |
| video_path_list = get_video_path_list( | |
| video_folder=args.video_folder, | |
| video_metadata_path=args.video_metadata_path, | |
| video_path_column=args.video_path_column | |
| ) | |
| 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() | |
| saved_video_path_list = [os.path.join(args.video_folder, video_path) for video_path in saved_video_path_list] | |
| video_path_list = list(set(video_path_list).difference(set(saved_video_path_list))) | |
| # Sorting to guarantee the same result for each process. | |
| video_path_list = natsorted(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.") | |
| if args.asethetic_score_metadata_path is not None: | |
| if args.asethetic_score_metadata_path.endswith(".csv"): | |
| asethetic_score_df = pd.read_csv(args.asethetic_score_metadata_path) | |
| elif args.asethetic_score_metadata_path.endswith(".jsonl"): | |
| asethetic_score_df = pd.read_json(args.asethetic_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( | |
| 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"] < args.asethetic_score_threshold] | |
| filtered_video_path_list = filtered_asethetic_score_df[args.video_path_column].tolist() | |
| filtered_video_path_list = [os.path.join(args.video_folder, video_path) for video_path in filtered_video_path_list] | |
| video_path_list = list(set(video_path_list).difference(set(filtered_video_path_list))) | |
| # Sorting to guarantee the same result for each process. | |
| video_path_list = natsorted(video_path_list) | |
| logger.info(f"Load {args.asethetic_score_metadata_path} and filter {len(filtered_video_path_list)} videos.") | |
| state = PartialState() | |
| ocr_reader = init_ocr_reader(device=state.device) | |
| # The workaround can be removed after https://github.com/huggingface/accelerate/pull/2781 is released. | |
| index = len(video_path_list) - len(video_path_list) % state.num_processes | |
| logger.info(f"Drop {len(video_path_list) % state.num_processes} videos to avoid duplicates in state.split_between_processes.") | |
| video_path_list = video_path_list[:index] | |
| result_list = [] | |
| with state.split_between_processes(video_path_list) as splitted_video_path_list: | |
| for i, video_path in enumerate(tqdm(splitted_video_path_list)): | |
| video_meta_info = compute_text_score(video_path, ocr_reader) | |
| result_list.append(video_meta_info) | |
| if i != 0 and i % args.saved_freq == 0: | |
| state.wait_for_everyone() | |
| gathered_result_list = gather_object(result_list) | |
| if state.is_main_process: | |
| 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") | |
| 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: | |
| logger.info(len(gathered_result_list)) | |
| if 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") | |
| logger.info(f"Save the final result to {args.saved_path}.") | |
| if __name__ == "__main__": | |
| main() |