yourComplete/quickstart-trackio

1import os2import shutil3import uuid4from abc import ABC, abstractmethod5from pathlib import Path6from typing import Literal7 8import numpy as np9from PIL import Image as PILImage10 11try:  # absolute imports when installed12    from trackio.file_storage import FileStorage13    from trackio.utils import MEDIA_DIR14    from trackio.video_writer import write_video15except ImportError:  # relative imports for local execution on Spaces16    from file_storage import FileStorage17    from utils import MEDIA_DIR18    from video_writer import write_video19 20 21class TrackioMedia(ABC):22    """23    Abstract base class for Trackio media objects24    Provides shared functionality for file handling and serialization.25    """26 27    TYPE: str28 29    def __init_subclass__(cls, **kwargs):30        """Ensure subclasses define the TYPE attribute."""31        super().__init_subclass__(**kwargs)32        if not hasattr(cls, "TYPE") or cls.TYPE is None:33            raise TypeError(f"Class {cls.__name__} must define TYPE attribute")34 35    def __init__(self, value, caption: str | None = None):36        self.caption = caption37        self._value = value38        self._file_path: Path | None = None39 40        # Validate file existence for string/Path inputs41        if isinstance(self._value, str | Path):42            if not os.path.isfile(self._value):43                raise ValueError(f"File not found: {self._value}")44 45    def _file_extension(self) -> str:46        if self._file_path:47            return self._file_path.suffix[1:].lower()48        if isinstance(self._value, str | Path):49            path = Path(self._value)50            return path.suffix[1:].lower()51        if hasattr(self, "_format") and self._format:52            return self._format53        return "unknown"54 55    def _get_relative_file_path(self) -> Path | None:56        return self._file_path57 58    def _get_absolute_file_path(self) -> Path | None:59        if self._file_path:60            return MEDIA_DIR / self._file_path61        return None62 63    def _save(self, project: str, run: str, step: int = 0):64        if self._file_path:65            return66 67        media_dir = FileStorage.init_project_media_path(project, run, step)68        filename = f"{uuid.uuid4()}.{self._file_extension()}"69        file_path = media_dir / filename70 71        # Delegate to subclass-specific save logic72        self._save_media(file_path)73 74        self._file_path = file_path.relative_to(MEDIA_DIR)75 76    @abstractmethod77    def _save_media(self, file_path: Path):78        """79        Performs the actual media saving logic.80        """81        pass82 83    def _to_dict(self) -> dict:84        if not self._file_path:85            raise ValueError("Media must be saved to file before serialization")86        return {87            "_type": self.TYPE,88            "file_path": str(self._get_relative_file_path()),89            "caption": self.caption,90        }91 92 93TrackioImageSourceType = str | Path | np.ndarray | PILImage.Image94 95 96class TrackioImage(TrackioMedia):97    """98    Initializes an Image object.99 100    Example:101        ```python102        import trackio103        import numpy as np104        from PIL import Image105 106        # Create an image from numpy array107        image_data = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)108        image = trackio.Image(image_data, caption="Random image")109        trackio.log({"my_image": image})110 111        # Create an image from PIL Image112        pil_image = Image.new('RGB', (100, 100), color='red')113        image = trackio.Image(pil_image, caption="Red square")114        trackio.log({"red_image": image})115 116        # Create an image from file path117        image = trackio.Image("path/to/image.jpg", caption="Photo from file")118        trackio.log({"file_image": image})119        ```120 121    Args:122        value (`str`, `Path`, `numpy.ndarray`, or `PIL.Image`, *optional*):123            A path to an image, a PIL Image, or a numpy array of shape (height, width, channels).124        caption (`str`, *optional*):125            A string caption for the image.126    """127 128    TYPE = "trackio.image"129 130    def __init__(self, value: TrackioImageSourceType, caption: str | None = None):131        super().__init__(value, caption)132        self._format: str | None = None133 134        if (135            isinstance(self._value, np.ndarray | PILImage.Image)136            and self._format is None137        ):138            self._format = "png"139 140    def _as_pil(self) -> PILImage.Image | None:141        try:142            if isinstance(self._value, np.ndarray):143                arr = np.asarray(self._value).astype("uint8")144                return PILImage.fromarray(arr).convert("RGBA")145            if isinstance(self._value, PILImage.Image):146                return self._value.convert("RGBA")147        except Exception as e:148            raise ValueError(f"Failed to process image data: {self._value}") from e149        return None150 151    def _save_media(self, file_path: Path):152        if pil := self._as_pil():153            pil.save(file_path, format=self._format)154        elif isinstance(self._value, str | Path):155            if os.path.isfile(self._value):156                shutil.copy(self._value, file_path)157            else:158                raise ValueError(f"File not found: {self._value}")159 160 161TrackioVideoSourceType = str | Path | np.ndarray162TrackioVideoFormatType = Literal["gif", "mp4", "webm"]163 164 165class TrackioVideo(TrackioMedia):166    """167    Initializes a Video object.168 169    Example:170        ```python171        import trackio172        import numpy as np173 174        # Create a simple video from numpy array175        frames = np.random.randint(0, 255, (10, 3, 64, 64), dtype=np.uint8)176        video = trackio.Video(frames, caption="Random video", fps=30)177 178        # Create a batch of videos179        batch_frames = np.random.randint(0, 255, (3, 10, 3, 64, 64), dtype=np.uint8)180        batch_video = trackio.Video(batch_frames, caption="Batch of videos", fps=15)181 182        # Create video from file path183        video = trackio.Video("path/to/video.mp4", caption="Video from file")184        ```185 186    Args:187        value (`str`, `Path`, or `numpy.ndarray`, *optional*):188            A path to a video file, or a numpy array.189            The array should be of type `np.uint8` with RGB values in the range `[0, 255]`.190            It is expected to have shape of either (frames, channels, height, width) or (batch, frames, channels, height, width).191            For the latter, the videos will be tiled into a grid.192        caption (`str`, *optional*):193            A string caption for the video.194        fps (`int`, *optional*):195            Frames per second for the video. Only used when value is an ndarray. Default is `24`.196        format (`Literal["gif", "mp4", "webm"]`, *optional*):197            Video format ("gif", "mp4", or "webm"). Only used when value is an ndarray. Default is "gif".198    """199 200    TYPE = "trackio.video"201 202    def __init__(203        self,204        value: TrackioVideoSourceType,205        caption: str | None = None,206        fps: int | None = None,207        format: TrackioVideoFormatType | None = None,208    ):209        super().__init__(value, caption)210        if isinstance(value, np.ndarray):211            if format is None:212                format = "gif"213            if fps is None:214                fps = 24215        self._fps = fps216        self._format = format217 218    @property219    def _codec(self) -> str:220        match self._format:221            case "gif":222                return "gif"223            case "mp4":224                return "h264"225            case "webm":226                return "vp9"227            case _:228                raise ValueError(f"Unsupported format: {self._format}")229 230    def _save_media(self, file_path: Path):231        if isinstance(self._value, np.ndarray):232            video = TrackioVideo._process_ndarray(self._value)233            write_video(file_path, video, fps=self._fps, codec=self._codec)234        elif isinstance(self._value, str | Path):235            if os.path.isfile(self._value):236                shutil.copy(self._value, file_path)237            else:238                raise ValueError(f"File not found: {self._value}")239 240    @staticmethod241    def _process_ndarray(value: np.ndarray) -> np.ndarray:242        # Verify value is either 4D (single video) or 5D array (batched videos).243        # Expected format: (frames, channels, height, width) or (batch, frames, channels, height, width)244        if value.ndim < 4:245            raise ValueError(246                "Video requires at least 4 dimensions (frames, channels, height, width)"247            )248        if value.ndim > 5:249            raise ValueError(250                "Videos can have at most 5 dimensions (batch, frames, channels, height, width)"251            )252        if value.ndim == 4:253            # Reshape to 5D with single batch: (1, frames, channels, height, width)254            value = value[np.newaxis, ...]255 256        value = TrackioVideo._tile_batched_videos(value)257        return value258 259    @staticmethod260    def _tile_batched_videos(video: np.ndarray) -> np.ndarray:261        """262        Tiles a batch of videos into a grid of videos.263 264        Input format: (batch, frames, channels, height, width) - original FCHW format265        Output format: (frames, total_height, total_width, channels)266        """267        batch_size, frames, channels, height, width = video.shape268 269        next_pow2 = 1 << (batch_size - 1).bit_length()270        if batch_size != next_pow2:271            pad_len = next_pow2 - batch_size272            pad_shape = (pad_len, frames, channels, height, width)273            padding = np.zeros(pad_shape, dtype=video.dtype)274            video = np.concatenate((video, padding), axis=0)275            batch_size = next_pow2276 277        n_rows = 1 << ((batch_size.bit_length() - 1) // 2)278        n_cols = batch_size // n_rows279 280        # Reshape to grid layout: (n_rows, n_cols, frames, channels, height, width)281        video = video.reshape(n_rows, n_cols, frames, channels, height, width)282 283        # Rearrange dimensions to (frames, total_height, total_width, channels)284        video = video.transpose(2, 0, 4, 1, 5, 3)285        video = video.reshape(frames, n_rows * height, n_cols * width, channels)286        return video287