media_loader.py

"""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,
    )