Spaces:
Running on L40S
Running on L40S
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| import re | |
| from typing import List, Tuple | |
| IMAGE_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { | |
| # Our desired 256 resolution is the one below (commented). | |
| # Desired: "256": {"1,1": (336, 336), "4,3": (384, 288), "3,4": (288, 384), "16,9": (448, 256), "9,16": (256, 448)}, | |
| "256": { | |
| "1,1": (256, 256), | |
| "4,3": (320, 256), | |
| "3,4": (256, 320), | |
| "16,9": (320, 192), | |
| "9,16": (192, 320), | |
| }, | |
| "480": {"1,1": (640, 640), "4,3": (736, 544), "3,4": (544, 736), "16,9": (832, 480), "9,16": (480, 832)}, | |
| # 704 resolutions are nicely divisible by 32 | |
| "704": {"1,1": (960, 960), "4,3": (1088, 832), "3,4": (832, 1088), "16,9": (1280, 704), "9,16": (704, 1280)}, | |
| "720": {"1,1": (960, 960), "4,3": (1104, 832), "3,4": (832, 1104), "16,9": (1280, 720), "9,16": (720, 1280)}, | |
| # 768 for arena.ai | |
| "768": {"1,1": (1024, 1024), "4,3": (1184, 880), "3,4": (880, 1184), "16,9": (1360, 768), "9,16": (768, 1360)}, | |
| "1080": {"1,1": (1440, 1440), "4,3": (1664, 1248), "3,4": (1248, 1664), "16,9": (1920, 1080), "9,16": (1080, 1920)}, | |
| "1280": {"1,1": (1712, 1712), "4,3": (1968, 1472), "3,4": (1472, 1968), "16,9": (2272, 1280), "9,16": (1280, 2272)}, | |
| "2048": { | |
| "1,1": (2728, 2728), | |
| "4,3": (3160, 2368), | |
| "3,4": (2368, 3160), | |
| "16,9": (3640, 2048), | |
| "9,16": (2048, 3640), | |
| }, | |
| "gt_2048": { | |
| "1,1": (5464, 5464), | |
| "4,3": (6304, 4728), | |
| "3,4": (4728, 6304), | |
| "16,9": (7280, 4096), | |
| "9,16": (4096, 7280), | |
| }, | |
| } | |
| VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { | |
| # Our desired 256 resolution is the one below (commented). | |
| # Desired: "256": {"1,1": (336, 336), "4,3": (384, 288), "3,4": (288, 384), "16,9": (448, 256), "9,16": (256, 448)}, | |
| "256": { | |
| "1,1": (256, 256), | |
| "4,3": (320, 256), | |
| "3,4": (256, 320), | |
| "16,9": (320, 192), | |
| "9,16": (192, 320), | |
| }, | |
| "480": {"1,1": (640, 640), "4,3": (736, 544), "3,4": (544, 736), "16,9": (832, 480), "9,16": (480, 832)}, | |
| # 704 resolutions are nicely divisible by 32 | |
| "704": {"1,1": (960, 960), "4,3": (1088, 832), "3,4": (832, 1088), "16,9": (1280, 704), "9,16": (704, 1280)}, | |
| "720": {"1,1": (960, 960), "4,3": (1104, 832), "3,4": (832, 1104), "16,9": (1280, 720), "9,16": (720, 1280)}, | |
| # 768 for arena.ai | |
| "768": {"1,1": (1024, 1024), "4,3": (1184, 880), "3,4": (880, 1184), "16,9": (1360, 768), "9,16": (768, 1360)}, | |
| "1080": {"1,1": (1440, 1440), "4,3": (1664, 1248), "3,4": (1248, 1664), "16,9": (1920, 1080), "9,16": (1080, 1920)}, | |
| "1280": {"1,1": (1712, 1712), "4,3": (1968, 1472), "3,4": (1472, 1968), "16,9": (2272, 1280), "9,16": (1280, 2272)}, | |
| "2048": { | |
| "1,1": (2728, 2728), | |
| "4,3": (3160, 2368), | |
| "3,4": (2368, 3160), | |
| "16,9": (3640, 2048), | |
| "9,16": (2048, 3640), | |
| }, | |
| "gt_2048": { | |
| "1,1": (5464, 5464), | |
| "4,3": (6304, 4728), | |
| "3,4": (4728, 6304), | |
| "16,9": (7280, 4096), | |
| "9,16": (4096, 7280), | |
| }, | |
| } | |
| def get_aspect_ratios_from_wdinfos(wdinfos: list[str]) -> list[str]: | |
| aspect_ratios = [] | |
| for wdinfo in wdinfos: | |
| aspect_ratio_match = re.search(r"aspect_ratio_(\d+_\d+)", wdinfo) | |
| aspect_ratios.append(aspect_ratio_match.group(1)) | |
| return aspect_ratios | |
| def get_wdinfos_w_aspect_ratio(wdinfos: list[str]) -> List[Tuple[str, str]]: | |
| aspect_ratios = get_aspect_ratios_from_wdinfos(wdinfos) | |
| # return a list of (wdinfo_path, aspect_ratio) pairs | |
| return [(wdinfo, aspect_ratio.replace("_", ",")) for wdinfo, aspect_ratio in zip(wdinfos, aspect_ratios)] | |
| def parse_frame_range_from_wdinfo(wdinfo: str) -> tuple[int, int] | None: | |
| """ | |
| Parse frame range from wdinfo path. | |
| Args: | |
| wdinfo: wdinfo path string containing frames_X_Y pattern | |
| Returns: | |
| Tuple of (min_frames, max_frames) if found, None otherwise | |
| Example: | |
| >>> parse_frame_range_from_wdinfo("wdinfo/v4/tv_drama/resolution_720/aspect_ratio_16_9/frames_300_400/wdinfo.json") | |
| (300, 400) | |
| """ | |
| match = re.search(r"frames_(\d+)_(\d+)", wdinfo) | |
| if match: | |
| return (int(match.group(1)), int(match.group(2))) | |
| return None | |
| def filter_wdinfos_by_frame_range( | |
| wdinfos: list[str], | |
| min_frames: int | None = None, | |
| max_frames: int | None = None, | |
| ) -> list[str]: | |
| """ | |
| Filter wdinfo files based on frame range. | |
| The frame range in wdinfo path (e.g., frames_300_400) represents videos | |
| with frames between those values. This function filters wdinfo files | |
| based on the wdinfo's upper bound (wdinfo_max): | |
| - min_frames is EXCLUSIVE: wdinfo_max must be > min_frames | |
| - max_frames is INCLUSIVE: wdinfo_max must be <= max_frames | |
| Args: | |
| wdinfos: List of wdinfo paths | |
| min_frames: Minimum number of frames (exclusive). If None, no lower bound. | |
| max_frames: Maximum number of frames (inclusive). If None, no upper bound. | |
| Returns: | |
| Filtered list of wdinfo paths | |
| Example: | |
| >>> wdinfos = [ | |
| ... "wdinfo/frames_400_500/wdinfo.json", | |
| ... "wdinfo/frames_500_600/wdinfo.json", | |
| ... "wdinfo/frames_600_700/wdinfo.json", | |
| ... ] | |
| >>> filter_wdinfos_by_frame_range(wdinfos, min_frames=500, max_frames=600) | |
| ['wdinfo/frames_500_600/wdinfo.json'] | |
| # frames_400_500 excluded because wdinfo_max (500) <= min_frames (500) | |
| # frames_500_600 included because wdinfo_max (600) > min_frames (500) AND <= max_frames (600) | |
| # frames_600_700 excluded because wdinfo_max (700) > max_frames (600) | |
| """ | |
| if min_frames is None and max_frames is None: | |
| return wdinfos | |
| filtered = [] | |
| for wdinfo in wdinfos: | |
| frame_range = parse_frame_range_from_wdinfo(wdinfo) | |
| if frame_range is None: | |
| # If no frame range in path, include by default | |
| filtered.append(wdinfo) | |
| continue | |
| wdinfo_min, wdinfo_max = frame_range | |
| # Filter based on wdinfo's upper bound (wdinfo_max): | |
| # - min_frames is exclusive: wdinfo_max must be > min_frames | |
| # - max_frames is inclusive: wdinfo_max must be <= max_frames | |
| include = True | |
| if min_frames is not None and wdinfo_max <= min_frames: | |
| include = False | |
| if max_frames is not None and wdinfo_max > max_frames: | |
| include = False | |
| if include: | |
| filtered.append(wdinfo) | |
| return filtered | |