"""Media loading utilities for ZipSplatPlus Space. Supports still images (JPEG, PNG, WebP, HEIC/HEIF), videos (MOV, MP4, M4V), Apple Live Photo pairs, and embedded video extraction from single HEIC files. """ from dataclasses import dataclass, field from pathlib import Path import tempfile from typing import List, Optional, Sequence, Tuple, Union try: import imageio.v2 as imageio except ImportError: import imageio import numpy as np import torch from PIL import Image, ImageOps IMAGE_EXTENSIONS = { ".jpg", ".jpeg", ".png", ".webp", ".heic", ".heif", } VIDEO_EXTENSIONS = { ".mov", ".mp4", ".m4v", } @dataclass class MediaLoadResult: images: List[torch.Tensor] image_count: int video_count: int decoded_video_frames: int selected_video_frames: int warnings: List[str] = field(default_factory=list) _HEIF_REGISTERED = False def register_heif() -> None: """Register Pillow HEIF plugin for decoding .heic and .heif images.""" global _HEIF_REGISTERED if not _HEIF_REGISTERED: try: from pillow_heif import register_heif_opener register_heif_opener() _HEIF_REGISTERED = True except Exception: pass register_heif() def to_tensor(image) -> torch.Tensor: """Convert HWC image (uint8 or float in [0, 1]) to (3, H, W) float in [0, 1].""" arr = np.asarray(image) arr = arr.astype(np.float32) / 255.0 if arr.dtype == np.uint8 else arr.astype(np.float32) tensor = torch.from_numpy(arr) if tensor.ndim == 3 and tensor.shape[-1] in (3, 4): tensor = tensor[..., :3].permute(2, 0, 1) return tensor.contiguous() def load_image(path: Union[Path, str]) -> torch.Tensor: """Load an image to a (3, H, W) float tensor in [0, 1].""" register_heif() path = Path(path) suffix = path.suffix.lower() if suffix in {".heic", ".heif"} and not _HEIF_REGISTERED: raise ValueError( f"Cannot decode HEIC image '{path.name}': pillow-heif is not registered/installed in Python environment." ) try: with Image.open(path) as img: img = ImageOps.exif_transpose(img) img = img.convert("RGB") return to_tensor(img) except Image.UnidentifiedImageError as e: if suffix in {".heic", ".heif"}: raise ValueError( f"UnidentifiedImageError: Could not decode HEIC image '{path.name}'. " "The file may be corrupted or use an unsupported HEIF container variant." ) from e raise ValueError( f"UnidentifiedImageError: Image format of '{path.name}' was not recognized by Pillow." ) from e except Exception as e: raise ValueError(f"Could not load image '{path.name}': {e}") from e def extract_embedded_video_from_heic(path: Union[Path, str]) -> Optional[Path]: """Attempt to extract an embedded MP4/MOV video stream from an Apple Live Photo HEIC container.""" try: path = Path(path) with open(path, "rb") as f: data = f.read() signatures = [b"ftypmp42", b"ftypisom", b"ftypqt ", b"ftypMSNV"] best_pos = -1 for sig in signatures: pos = data.find(sig, 12) # search after initial ftyp box if pos >= 4: if best_pos == -1 or pos < best_pos: best_pos = pos if best_pos >= 4: start = best_pos - 4 video_data = data[start:] if len(video_data) > 1000: tmp_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) tmp_video.write(video_data) tmp_video.close() return Path(tmp_video.name) except Exception: pass return None def load_video( path: Union[Path, str], num_frames: Optional[int] = 8, stride: Optional[int] = None ) -> List[torch.Tensor]: """Load video frames as a list of (3, H, W) float tensors in [0, 1].""" reader = imageio.get_reader(str(path)) try: total = reader.count_frames() if stride is not None: indices = list(range(0, total, stride)) else: n = min(num_frames, total) indices = torch.linspace(0, total - 1, n).round().long().tolist() frames = [to_tensor(reader.get_data(i)) for i in indices] except Exception: frames = [to_tensor(f) for f in reader] if stride is not None: frames = frames[::stride] elif num_frames is not None and len(frames) > num_frames: idx = torch.linspace(0, len(frames) - 1, num_frames).round().long().tolist() frames = [frames[i] for i in idx] finally: reader.close() return frames def _is_duplicate_frame( candidate: torch.Tensor, existing: List[torch.Tensor], threshold: float = 0.005 ) -> bool: """Check if candidate frame is identical or nearly identical to any existing frame.""" for ex in existing: if candidate.shape == ex.shape: if torch.equal(candidate, ex): return True diff = (candidate - ex).abs().mean().item() if diff < threshold: return True return False def load_media_views( paths: Sequence[Union[str, Path]], *, max_views: int = 24, frames_per_video: int = 8, ) -> MediaLoadResult: """Load media files (images, videos, Live Photo pairs, single HEIC with embedded video) into a bounded set of view tensors.""" warnings: List[str] = [] normalized_paths: List[Path] = [Path(p) for p in paths] valid_images: List[Path] = [] valid_videos: List[Path] = [] for p in normalized_paths: suffix = p.suffix.lower() if suffix in IMAGE_EXTENSIONS: valid_images.append(p) elif suffix in VIDEO_EXTENSIONS: valid_videos.append(p) else: warnings.append(f"Skipped unsupported file: {p.name}") image_stems = {p.stem: p for p in valid_images} paired_videos: List[Path] = [] independent_videos: List[Path] = [] for vp in valid_videos: if vp.stem in image_stems: paired_videos.append(vp) else: independent_videos.append(vp) paired_video_stems = {vp.stem for vp in paired_videos} ordered_videos = paired_videos + independent_videos selected_views: List[torch.Tensor] = [] image_count = 0 video_count = 0 decoded_video_frames = 0 selected_video_frames = 0 # 1. Process still images (and check for embedded Live Photo videos in standalone HEIC files) for p in valid_images: if len(selected_views) >= max_views: warnings.append(f"Maximum view limit ({max_views}) reached; skipped image {p.name}.") continue try: tensor = load_image(p) if not _is_duplicate_frame(tensor, selected_views): selected_views.append(tensor) image_count += 1 # Auto-extract embedded Live Photo video if available and no separate paired video was uploaded if p.suffix.lower() in {".heic", ".heif"} and p.stem not in paired_video_stems: embedded_video_path = extract_embedded_video_from_heic(p) if embedded_video_path: try: frames = load_video(embedded_video_path, num_frames=frames_per_video) video_count += 1 decoded_video_frames += len(frames) for frame in frames: if len(selected_views) >= max_views: break if not _is_duplicate_frame(frame, selected_views): selected_views.append(frame) selected_video_frames += 1 finally: try: embedded_video_path.unlink(missing_ok=True) except Exception: pass except Exception as e: warnings.append(f"Skipped corrupt or unreadable image {p.name}: {e}") # 2. Process uploaded video files for vp in ordered_videos: if len(selected_views) >= max_views: warnings.append(f"Maximum view limit ({max_views}) reached; skipped video {vp.name}.") continue try: frames = load_video(vp, num_frames=frames_per_video) video_count += 1 decoded_video_frames += len(frames) added_for_video = 0 for frame in frames: if len(selected_views) >= max_views: break if not _is_duplicate_frame(frame, selected_views): selected_views.append(frame) selected_video_frames += 1 added_for_video += 1 if added_for_video < len(frames) and len(selected_views) >= max_views: warnings.append( f"Capped video frames from {vp.name} due to max_views limit ({max_views})." ) except Exception as e: warnings.append(f"Skipped corrupt or unreadable video {vp.name}: {e}") if not selected_views: msg = "No usable image or video views could be loaded from the provided inputs." if warnings: msg += "\nWarnings:\n" + "\n".join(warnings) raise ValueError(msg) return MediaLoadResult( images=selected_views, image_count=image_count, video_count=video_count, decoded_video_frames=decoded_video_frames, selected_video_frames=selected_video_frames, warnings=warnings, )