Spaces:
Running on Zero
Running on Zero
| """Input handling: a directory of images or a video file -> list of frame paths.""" | |
| from __future__ import annotations | |
| import glob | |
| import os | |
| import tempfile | |
| IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp") | |
| VIDEO_EXTS = (".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v") | |
| def _natural_key(path: str): | |
| import re | |
| name = os.path.basename(path) | |
| return [int(c) if c.isdigit() else c.lower() for c in re.split(r"(\d+)", name)] | |
| def load_frame_paths(input_path: str, max_frames: int | None = None, | |
| stride: int = 1, work_dir: str | None = None): | |
| """Return a sorted list of image paths from a folder or extracted from a video. | |
| Args: | |
| input_path: a directory of images or a single video file. | |
| max_frames: keep at most this many frames (after stride). | |
| stride: take every ``stride``-th frame. | |
| work_dir: where to extract video frames (a temp dir if ``None``). | |
| Returns: | |
| (frame_paths, fps) where ``fps`` is the source video fps or ``None``. | |
| """ | |
| if os.path.isdir(input_path): | |
| paths = [] | |
| for ext in IMAGE_EXTS: | |
| paths += glob.glob(os.path.join(input_path, f"*{ext}")) | |
| paths += glob.glob(os.path.join(input_path, f"*{ext.upper()}")) | |
| paths = sorted(set(paths), key=_natural_key) | |
| if not paths: | |
| raise FileNotFoundError(f"No images found in directory: {input_path}") | |
| paths = paths[::stride] | |
| if max_frames: | |
| paths = paths[:max_frames] | |
| return paths, None | |
| ext = os.path.splitext(input_path)[1].lower() | |
| if ext in VIDEO_EXTS: | |
| return _extract_video_frames(input_path, max_frames, stride, work_dir) | |
| if ext in IMAGE_EXTS: # a single image | |
| return [input_path], None | |
| raise ValueError(f"Unsupported input '{input_path}'. Provide a single image, " | |
| f"an image folder, or a video file ({', '.join(VIDEO_EXTS)}).") | |
| def _extract_video_frames(video_path, max_frames, stride, work_dir): | |
| import imageio.v2 as imageio | |
| from PIL import Image | |
| if work_dir is None: | |
| work_dir = tempfile.mkdtemp(prefix="faceanything_frames_") | |
| os.makedirs(work_dir, exist_ok=True) | |
| reader = imageio.get_reader(video_path) | |
| try: | |
| fps = float(reader.get_meta_data().get("fps", 0)) or None | |
| except Exception: | |
| fps = None | |
| paths = [] | |
| count = 0 | |
| for i, frame in enumerate(reader): | |
| if i % stride != 0: | |
| continue | |
| out = os.path.join(work_dir, f"frame_{count:06d}.png") | |
| Image.fromarray(frame).save(out) | |
| paths.append(out) | |
| count += 1 | |
| if max_frames and count >= max_frames: | |
| break | |
| reader.close() | |
| if not paths: | |
| raise RuntimeError(f"No frames extracted from video: {video_path}") | |
| return paths, fps | |
| def find_mask_paths(frame_paths, mask_dir: str | None): | |
| """Best-effort match of a foreground mask for each frame. | |
| Tries an explicit ``mask_dir`` (matching by basename), then a sibling | |
| ``masks/`` directory next to the frames. Returns a list aligned with | |
| ``frame_paths`` (entries may be ``None``). | |
| """ | |
| candidates = [] | |
| if mask_dir and os.path.isdir(mask_dir): | |
| candidates.append(mask_dir) | |
| if frame_paths: | |
| sibling = os.path.join(os.path.dirname(os.path.dirname(frame_paths[0])), "masks") | |
| if os.path.isdir(sibling): | |
| candidates.append(sibling) | |
| masks = [] | |
| for fp in frame_paths: | |
| stem = os.path.splitext(os.path.basename(fp))[0] | |
| found = None | |
| for d in candidates: | |
| for ext in IMAGE_EXTS: | |
| for cand in (os.path.join(d, stem + ext), | |
| os.path.join(d, os.path.basename(fp))): | |
| if os.path.exists(cand): | |
| found = cand | |
| break | |
| if found: | |
| break | |
| if found: | |
| break | |
| masks.append(found) | |
| return masks | |