Add video tracking and Media tab visualization

aim/sdk/__init__.py

@@ -1,5 +1,5 @@
 # pre-defined sequences and custom objects
-from aim.sdk.objects import Audio, Distribution, Figure, Image, Text
+from aim.sdk.objects import Audio, Distribution, Figure, Image, Text, Video
 from aim.sdk.repo import Repo
 
 # SDK aliases
@@ -10,4 +10,5 @@
 from aim.sdk.sequences.image_sequence import Images
 from aim.sdk.sequences.metric import Metric
 from aim.sdk.sequences.text_sequence import Texts
+from aim.sdk.sequences.video_sequence import Videos
 from aim.sdk.training_flow import TrainingFlow

aim/sdk/objects/__init__.py

@@ -3,3 +3,4 @@
 from aim.sdk.objects.figure import Figure
 from aim.sdk.objects.image import Image
 from aim.sdk.objects.text import Text
+from aim.sdk.objects.video import Video

aim/sdk/objects/video.py

@@ -0,0 +1,131 @@
+import io
+import os.path
+
+from aim.storage.inmemorytreeview import InMemoryTreeView
+from aim.storage.object import CustomObject
+from aim.storage.types import BLOB
+
+
+@CustomObject.alias('aim.video')
+class Video(CustomObject):
+    """Video object used to store video files in Aim repositories.
+
+    Args:
+         path (:obj:`str`, optional): Video file path. Path-backed videos are
+             read when the run tracking worker encodes the object, which keeps
+             ``run.track()`` calls cheap for large videos when async tracking is
+             enabled.
+         data (:obj:`bytes` or :obj:`io.BytesIO`, optional): Video bytes.
+         fps (:obj:`float`, optional): Video frame rate.
+         format (:obj:`str`, optional): Video format. Inferred from ``path``
+             when possible. Supported formats are ``mp4``, ``m4v``, ``gif``,
+             ``mov`` and ``webm``.
+         caption (:obj:`str`, optional): Optional video caption. '' by default.
+    """
+
+    AIM_NAME = 'aim.video'
+
+    MP4 = 'mp4'
+    M4V = 'm4v'
+    GIF = 'gif'
+    MOV = 'mov'
+    WEBM = 'webm'
+
+    video_formats = (MP4, M4V, GIF, MOV, WEBM)
+
+    def __init__(self, path: str = None, *, data=None, fps: float = None, format: str = None, caption: str = ''):
+        super().__init__()
+
+        if path is None and data is None:
+            raise ValueError('Either video path or data must be provided.')
+
+        if path is not None:
+            if not os.path.exists(path) or not os.path.isfile(path):
+                raise ValueError('Invalid video file path.')
+            if format is None:
+                format = os.path.splitext(path)[1].lower().lstrip('.')
+        elif isinstance(data, io.BytesIO):
+            data = data.read()
+
+        if path is None and not isinstance(data, bytes):
+            raise TypeError('Content is not a byte-stream object.')
+
+        video_format = (format or '').lower()
+        if video_format not in self.video_formats:
+            raise ValueError(f'Invalid video format is provided. Must be one of {self.video_formats}')
+
+        self.storage['caption'] = caption
+        self.storage['format'] = video_format
+        self.storage['fps'] = fps
+        if path is not None:
+            self.storage['source_path'] = os.path.abspath(path)
+            self.storage['size'] = os.path.getsize(path)
+        else:
+            self.storage['size'] = len(data)
+            self.storage['data'] = BLOB(data=data)
+
+    @property
+    def caption(self) -> str:
+        return self.storage['caption']
+
+    @property
+    def format(self) -> str:
+        return self.storage['format']
+
+    @property
+    def fps(self):
+        return self.storage['fps']
+
+    @property
+    def size(self) -> int:
+        return self.storage.get('size', 0)
+
+    def json(self):
+        """Dump video metadata to a dict."""
+        return {
+            'caption': self.caption,
+            'format': self.format,
+            'fps': self.fps,
+            'size': self.size,
+        }
+
+    def __deepcopy__(self, memodict=None):
+        if memodict is None:
+            memodict = {}
+
+        storage = InMemoryTreeView(container={})
+        for key in ('caption', 'format', 'fps', 'size', 'source_path', 'data'):
+            try:
+                storage[key] = self.storage[key]
+            except KeyError:
+                pass
+        result = self.__class__.__new__(self.__class__, _storage=storage)
+        memodict[id(self)] = result
+        return result
+
+    def get(self) -> io.BytesIO:
+        """Return video bytes as an in-memory buffer."""
+        try:
+            bs = self.storage['data']
+            return io.BytesIO(bytes(bs))
+        except KeyError:
+            pass
+        source_path = self.storage.get('source_path')
+        if source_path:
+            with open(source_path, 'rb') as fs:
+                return io.BytesIO(fs.read())
+        return io.BytesIO()
+
+    def _aim_encode(self):
+        try:
+            self.storage['data']
+        except KeyError:
+            source_path = self.storage.get('source_path')
+            if not source_path:
+                raise ValueError('Video data is missing.')
+            with open(source_path, 'rb') as fs:
+                self.storage['data'] = BLOB(data=fs.read())
+            self.storage['size'] = os.path.getsize(source_path)
+        if self.storage.get('source_path'):
+            del self.storage['source_path']
+        return self.AIM_NAME, self.storage[...]

aim/sdk/repo.py

@@ -630,6 +630,15 @@ def query_audios(
 
         return QuerySequenceCollection(repo=self, seq_cls=Audios, query=query, report_mode=report_mode)
 
+    def query_videos(
+        self, query: str = '', report_mode: QueryReportMode = QueryReportMode.PROGRESS_BAR
+    ) -> QuerySequenceCollection:
+        """Get video collections satisfying query expression."""
+        self._prepare_runs_cache()
+        from aim.sdk.sequences.video_sequence import Videos
+
+        return QuerySequenceCollection(repo=self, seq_cls=Videos, query=query, report_mode=report_mode)
+
     def query_figure_objects(
         self, query: str = '', report_mode: QueryReportMode = QueryReportMode.PROGRESS_BAR
     ) -> QuerySequenceCollection:

aim/sdk/run.py

@@ -50,6 +50,7 @@
     from aim.sdk.sequences.image_sequence import Images
     from aim.sdk.sequences.metric import Metric
     from aim.sdk.sequences.text_sequence import Texts
+    from aim.sdk.sequences.video_sequence import Videos
     from pandas import DataFrame
 
 logger = logging.getLogger(__name__)
@@ -606,6 +607,10 @@ def get_audio_sequence(self, name: str, context: Context) -> Optional['Audios']:
         """
         return self._get_sequence('audios', name, context)
 
+    def get_video_sequence(self, name: str, context: Context) -> Optional['Videos']:
+        """Retrieve videos sequence by its name and context."""
+        return self._get_sequence('videos', name, context)
+
     def get_distribution_sequence(self, name: str, context: Context) -> Optional['Distributions']:
         """Retrieve distributions sequence by it's name and context.
 

aim/sdk/sequences/sequence_type_map.py

@@ -7,6 +7,8 @@
     'list(aim.image)': 'images',
     'aim.audio': 'audios',
     'list(aim.audio)': 'audios',
+    'aim.video': 'videos',
+    'list(aim.video)': 'videos',
     'aim.text': 'texts',
     'list(aim.text)': 'texts',
     'aim.distribution': 'distributions',

aim/sdk/sequences/video_sequence.py

@@ -0,0 +1,17 @@
+from typing import Tuple, Union
+
+from aim.sdk.objects import Video
+from aim.sdk.sequence import MediaSequenceBase
+
+
+class Videos(MediaSequenceBase):
+    """Class representing series of Video objects or Video lists."""
+
+    @classmethod
+    def allowed_dtypes(cls) -> Union[str, Tuple[str, ...]]:
+        typename = Video.get_typename()
+        return typename, f'list({typename})'
+
+    @classmethod
+    def sequence_name(cls) -> str:
+        return 'videos'

aim/sdk/tracker.py

@@ -101,9 +101,11 @@ def __call__(
 
         if self._non_blocking:
             val = deepcopy(value)
-            self.repo.tracking_queue.register_task(
+            warn_queue_full = self.repo.tracking_queue.register_task(
                 self._track, val, track_time, name, step, epoch, context=context
-            ) or self.track_rate_warn()
+            )
+            if warn_queue_full:
+                self.track_rate_warn()
         else:
             self._track(value, track_time, name, step, epoch, context=context)
 

aim/storage/query.py

@@ -137,7 +137,7 @@ def query_add_default_expr(query: str) -> str:
 class RestrictedPythonQuery(Query):
     __slots__ = ('_checker', 'run_metadata_cache')
 
-    allowed_params = {'run', 'metric', 'images', 'audios', 'distributions', 'figures', 'texts'}
+    allowed_params = {'run', 'metric', 'images', 'audios', 'videos', 'distributions', 'figures', 'texts'}
 
     def __init__(self, query: str):
         stripped_query = strip_query(query)

aim/web/api/runs/object_views.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional
 
-from aim import Audios, Distributions, Figures, Images, Texts
+from aim import Audios, Distributions, Figures, Images, Texts, Videos
 from aim.sdk.sequence import Sequence
 from aim.sdk.sequence_collection import QuerySequenceCollection
 from aim.sdk.types import QueryReportMode
@@ -16,6 +16,7 @@
     RunTracesBatchApiIn,
     TextList,
     URIBatchIn,
+    VideoList,
 )
 from aim.web.api.runs.utils import (
     checked_query,
@@ -175,6 +176,12 @@ class AudioApiConfig(CustomObjectApiConfig):
     model = AudioList
 
 
+class VideoApiConfig(CustomObjectApiConfig):
+    sequence_type = Videos
+    resolve_blobs = False
+    model = VideoList
+
+
 class FigureApiConfig(CustomObjectApiConfig):
     sequence_type = Figures
     resolve_blobs = True

aim/web/api/runs/pydantic_models.py

@@ -221,6 +221,15 @@ class AudioInfo(BaseModel):
     index: int
 
 
+class VideoInfo(BaseModel):
+    caption: str
+    format: str
+    fps: Optional[float] = None
+    size: Optional[int] = None
+    blob_uri: str
+    index: int
+
+
 class DistributionInfo(BaseModel):
     data: EncodedNumpyArray
     bin_count: int
@@ -238,3 +247,4 @@ class NoteIn(BaseModel):
 ImageList = List[ImageInfo]
 TextList = List[TextInfo]
 AudioList = List[AudioInfo]
+VideoList = List[VideoInfo]

aim/web/api/runs/views.py

@@ -7,6 +7,7 @@
     FigureApiConfig,
     ImageApiConfig,
     TextApiConfig,
+    VideoApiConfig,
 )
 from aim.web.api.runs.pydantic_models import (
     MetricAlignApiIn,
@@ -374,4 +375,5 @@ def add_api_routes():
     TextApiConfig.register_endpoints(runs_router)
     DistributionApiConfig.register_endpoints(runs_router)
     AudioApiConfig.register_endpoints(runs_router)
+    VideoApiConfig.register_endpoints(runs_router)
     FigureApiConfig.register_endpoints(runs_router)

aim/web/ui/src/config/analytics/analyticsKeysMap.ts

@@ -322,6 +322,11 @@ export const ANALYTICS_EVENT_KEYS = {
         clickApplyButton: '[RunDetail] [Audios] Click apply button',
         changeContext: '[RunDetail] [Audios] Change context',
       },
+      videos: {
+        tabView: '[RunDetail] [Videos] Tab view',
+        clickApplyButton: '[RunDetail] [Videos] Click apply button',
+        changeContext: '[RunDetail] [Videos] Change context',
+      },
       figures: {
         tabView: '[RunDetail] [Figures] Tab view',
         clickApplyButton: '[RunDetail] [Figures] Click apply button',
@@ -370,5 +375,6 @@ export const ANALYTICS_EVENT_KEYS = {
   },
   figures: {} as any,
   audios: {} as any,
+  videos: {} as any,
   text: {} as any,
 };