diff --git a/aim/sdk/__init__.py b/aim/sdk/__init__.py index f7c190da1a..cd1441bea9 100644 --- a/aim/sdk/__init__.py +++ b/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 diff --git a/aim/sdk/objects/__init__.py b/aim/sdk/objects/__init__.py index c710bcdba7..0e455bda15 100644 --- a/aim/sdk/objects/__init__.py +++ b/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 diff --git a/aim/sdk/objects/video.py b/aim/sdk/objects/video.py new file mode 100644 index 0000000000..29e9b41eff --- /dev/null +++ b/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[...] diff --git a/aim/sdk/repo.py b/aim/sdk/repo.py index 1ffef1c9b4..c71e2bf96a 100644 --- a/aim/sdk/repo.py +++ b/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: diff --git a/aim/sdk/run.py b/aim/sdk/run.py index b53bdf72af..8480379425 100644 --- a/aim/sdk/run.py +++ b/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. diff --git a/aim/sdk/sequences/sequence_type_map.py b/aim/sdk/sequences/sequence_type_map.py index 5f7476fed3..e6b255adbc 100644 --- a/aim/sdk/sequences/sequence_type_map.py +++ b/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', diff --git a/aim/sdk/sequences/video_sequence.py b/aim/sdk/sequences/video_sequence.py new file mode 100644 index 0000000000..c7729e2269 --- /dev/null +++ b/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' diff --git a/aim/sdk/tracker.py b/aim/sdk/tracker.py index 1fd36cd3e2..e19742420a 100644 --- a/aim/sdk/tracker.py +++ b/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) diff --git a/aim/storage/query.py b/aim/storage/query.py index 82de23657a..527b651468 100644 --- a/aim/storage/query.py +++ b/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) diff --git a/aim/web/api/runs/object_views.py b/aim/web/api/runs/object_views.py index 6a4b4d0e93..db3d5080a3 100644 --- a/aim/web/api/runs/object_views.py +++ b/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 diff --git a/aim/web/api/runs/pydantic_models.py b/aim/web/api/runs/pydantic_models.py index d27e9a337a..3506e96dd3 100644 --- a/aim/web/api/runs/pydantic_models.py +++ b/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] diff --git a/aim/web/api/runs/views.py b/aim/web/api/runs/views.py index 014cd01796..a78107f4fe 100644 --- a/aim/web/api/runs/views.py +++ b/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) diff --git a/aim/web/ui/src/config/analytics/analyticsKeysMap.ts b/aim/web/ui/src/config/analytics/analyticsKeysMap.ts index 475b4f2ee5..e11fd565ac 100644 --- a/aim/web/ui/src/config/analytics/analyticsKeysMap.ts +++ b/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, }; diff --git a/aim/web/ui/src/pages/RunDetail/RunDetail.tsx b/aim/web/ui/src/pages/RunDetail/RunDetail.tsx index 3822c7cb24..095fb16012 100644 --- a/aim/web/ui/src/pages/RunDetail/RunDetail.tsx +++ b/aim/web/ui/src/pages/RunDetail/RunDetail.tsx @@ -64,6 +64,10 @@ const RunDetailMetricsAndSystemTab = React.lazy( /* webpackChunkName: "RunDetailMetricsAndSystemTab" */ './RunDetailMetricsAndSystemTab' ), ); +const RunDetailMediaTab = React.lazy( + () => + import(/* webpackChunkName: "RunDetailMediaTab" */ './RunDetailMediaTab'), +); const TraceVisualizationContainer = React.lazy( () => import( @@ -89,13 +93,20 @@ const tabs: Record = { metrics: 'Metrics', system: 'System', distributions: 'Distributions', - images: 'Images', - audios: 'Audios', + media: 'Media', texts: 'Texts', figures: 'Figures', settings: 'Settings', }; +const mediaRouteKeys = ['images', 'videos', 'audios']; + +function isMediaRouteKey( + value?: string, +): value is 'images' | 'videos' | 'audios' { + return value === 'images' || value === 'videos' || value === 'audios'; +} + function RunDetail(): React.FunctionComponentElement { const { runHash } = useParams<{ runHash: string }>(); const history = useHistory(); @@ -113,13 +124,24 @@ function RunDetail(): React.FunctionComponentElement { function redirect(): void { const splitPathname: string[] = pathname.split('/'); const path: string = `${url}/overview`; - if (splitPathname.length > 4) { + if (splitPathname.length > 5) { + history.replace(path); + setActiveTab(path); + return; + } + if ( + splitPathname.length === 5 && + (splitPathname[3] !== 'media' || !isMediaRouteKey(splitPathname[4])) + ) { history.replace(path); setActiveTab(path); return; } if (splitPathname[3]) { - if (!Object.keys(tabs).includes(splitPathname[3])) { + if ( + !Object.keys(tabs).includes(splitPathname[3]) && + !isMediaRouteKey(splitPathname[3]) + ) { history.replace(path); } } else { @@ -175,18 +197,10 @@ function RunDetail(): React.FunctionComponentElement { traceInfo={runData?.runTraces} /> ), - images: ( - - ), - audios: ( - @@ -238,10 +252,29 @@ function RunDetail(): React.FunctionComponentElement { } function getCurrentTabValue(pathname: string, url: string) { + const splitPathname = pathname.split('/'); + if (splitPathname[3] === 'media' || isMediaRouteKey(splitPathname[3])) { + return `${url}/media`; + } const values = Object.keys(tabs).map((tabKey) => `${url}/${tabKey}`); return values.indexOf(pathname) === -1 ? false : pathname; } + function getTabContent(tabKey: string) { + if (isMediaRouteKey(tabKey)) { + return ( + + ); + } + return tabContent[tabKey]; + } + React.useEffect(() => { if (pathname !== activeTab) { setActiveTab(pathname); @@ -450,24 +483,12 @@ function RunDetail(): React.FunctionComponentElement { height='calc(100vh - 98px)' > - {Object.keys(tabs).map((tabKey: string) => ( - - - {tabKey === 'overview' ? ( -
- - -
- } - > - {tabContent[tabKey]} - - - ) : ( -
-
+ {[...Object.keys(tabs), ...mediaRouteKeys].map( + (tabKey: string) => ( + + + {tabKey === 'overview' ? ( +
@@ -478,11 +499,25 @@ function RunDetail(): React.FunctionComponentElement { {tabContent[tabKey]}
-
- )} - - - ))} + ) : ( +
+
+ + +
+ } + > + {getTabContent(tabKey)} + +
+
+ )} + + + ), + )} diff --git a/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.scss b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.scss new file mode 100644 index 0000000000..c61e65cc0d --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.scss @@ -0,0 +1,37 @@ +@use 'src/styles/abstracts' as *; + +.RunDetailMediaTab { + display: flex; + flex-direction: column; + height: 100%; + min-height: 0; + + &__tabs { + flex: 0 0 auto; + min-height: toRem(42px); + border: toRem(1px) solid $grayish; + border-bottom: none; + border-top-left-radius: toRem(6px); + border-top-right-radius: toRem(6px); + background: $white; + + .MuiTab-root { + min-height: toRem(42px); + } + } + + &__content { + flex: 1 1 auto; + min-height: 0; + + .TraceVisualizationWrapper { + .Menu { + border-top-left-radius: 0; + } + + .VisualizerArea { + border-top-right-radius: 0; + } + } + } +} diff --git a/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.tsx b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.tsx new file mode 100644 index 0000000000..303ec87023 --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/RunDetailMediaTab.tsx @@ -0,0 +1,80 @@ +import React from 'react'; +import { Link, useLocation } from 'react-router-dom'; + +import { Tab, Tabs } from '@material-ui/core'; + +import { TraceType } from 'services/models/runs/types'; + +import TraceVisualizationContainer from '../TraceVisualizationContainer'; +import { ITraceVisualizationContainerProps } from '../types'; + +import './RunDetailMediaTab.scss'; + +type MediaTraceType = Extract; + +type RunDetailMediaTabProps = Pick< + ITraceVisualizationContainerProps, + 'runHash' | 'traceInfo' | 'runParams' +> & { + basePath: string; + initialTraceType?: MediaTraceType; +}; + +const mediaTabs: { label: string; value: MediaTraceType }[] = [ + { label: 'Images', value: 'images' }, + { label: 'Videos', value: 'videos' }, + { label: 'Audio', value: 'audios' }, +]; + +function isMediaTraceType(value?: string): value is MediaTraceType { + return value === 'images' || value === 'videos' || value === 'audios'; +} + +function RunDetailMediaTab({ + basePath, + initialTraceType = 'images', + runHash, + runParams, + traceInfo, +}: RunDetailMediaTabProps): React.FunctionComponentElement { + const { pathname } = useLocation(); + const pathParts = pathname.split('/'); + const pathTraceType = pathParts[3] === 'media' ? pathParts[4] : pathParts[3]; + const activeTraceType = isMediaTraceType(pathTraceType) + ? pathTraceType + : initialTraceType; + + return ( +
+ + {mediaTabs.map((tab) => ( + + ))} + +
+ +
+
+ ); +} + +RunDetailMediaTab.displayName = 'RunDetailMediaTab'; + +export default React.memo(RunDetailMediaTab); diff --git a/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/index.ts b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/index.ts new file mode 100644 index 0000000000..f86b757bc5 --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/RunDetailMediaTab/index.ts @@ -0,0 +1,3 @@ +import RunDetailMediaTab from './RunDetailMediaTab'; + +export default RunDetailMediaTab; diff --git a/aim/web/ui/src/pages/RunDetail/RunOverviewTab/components/RunOverviewSidebar/RunOverviewSidebar.tsx b/aim/web/ui/src/pages/RunDetail/RunOverviewTab/components/RunOverviewSidebar/RunOverviewSidebar.tsx index ad1b5ce931..16eb1a2a00 100644 --- a/aim/web/ui/src/pages/RunDetail/RunOverviewTab/components/RunOverviewSidebar/RunOverviewSidebar.tsx +++ b/aim/web/ui/src/pages/RunDetail/RunOverviewTab/components/RunOverviewSidebar/RunOverviewSidebar.tsx @@ -77,14 +77,12 @@ function RunOverviewSidebar({ value: traces?.distributions?.length || 0, }, { - name: 'Images', - path: `${path}/images`, - value: traces?.images?.length || 0, - }, - { - name: 'Audios', - path: `${path}/audios`, - value: traces?.audios?.length || 0, + name: 'Media', + path: `${path}/media`, + value: + (traces?.images?.length || 0) + + (traces?.videos?.length || 0) + + (traces?.audios?.length || 0), }, { name: 'Texts', diff --git a/aim/web/ui/src/pages/RunDetail/TraceVisualizationContainer/TraceVisualizationContainer.tsx b/aim/web/ui/src/pages/RunDetail/TraceVisualizationContainer/TraceVisualizationContainer.tsx index 223bb24bd2..4f1d740c48 100644 --- a/aim/web/ui/src/pages/RunDetail/TraceVisualizationContainer/TraceVisualizationContainer.tsx +++ b/aim/web/ui/src/pages/RunDetail/TraceVisualizationContainer/TraceVisualizationContainer.tsx @@ -40,12 +40,16 @@ const AudiosVisualizer = React.lazy( () => import(/* webpackChunkName: "AudiosVisualizer" */ '../AudiosVisualizer'), ); +const VideosVisualizer = React.lazy( + () => + import(/* webpackChunkName: "VideosVisualizer" */ '../VideosVisualizer'), +); const traceTypeVisualization = { images: ImagesVisualizer, distributions: DistributionsVisualizer, audios: AudiosVisualizer, - videos: () => null, // @TODO add tracking event keys in analyticsKeysMap object + videos: VideosVisualizer, texts: TextsVisualizer, figures: PlotlyVisualizer, }; diff --git a/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.scss b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.scss new file mode 100644 index 0000000000..07cee8db48 --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.scss @@ -0,0 +1,60 @@ +@use 'src/styles/abstracts' as *; + +.VideosVisualizer { + height: calc(100% - 2.75rem); + overflow: auto; + padding: toRem(16px); + + &__grid { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(toRem(280px), 1fr)); + gap: toRem(16px); + align-items: start; + } + + &__card { + min-width: 0; + border: toRem(1px) solid $grayish; + border-radius: toRem(6px); + background: $white; + overflow: hidden; + } + + &__mediaBox { + position: relative; + width: 100%; + aspect-ratio: 16 / 9; + background: #101828; + + video, + img { + display: block; + width: 100%; + height: 100%; + object-fit: contain; + background: #101828; + } + } + + &__placeholder { + position: absolute; + inset: 0; + display: flex; + align-items: center; + justify-content: center; + background: #101828; + } + + &__meta { + display: flex; + flex-direction: column; + gap: toRem(4px); + padding: toRem(10px) toRem(12px) toRem(12px); + + h3 { + margin: 0; + line-height: 1.35; + overflow-wrap: anywhere; + } + } +} diff --git a/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.tsx b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.tsx new file mode 100644 index 0000000000..ec3003615e --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/VideosVisualizer.tsx @@ -0,0 +1,348 @@ +import React from 'react'; + +import BusyLoaderWrapper from 'components/BusyLoaderWrapper/BusyLoaderWrapper'; +import ErrorBoundary from 'components/ErrorBoundary/ErrorBoundary'; +import IllustrationBlock from 'components/IllustrationBlock/IllustrationBlock'; +import { Text } from 'components/kit'; + +import { IllustrationsEnum } from 'config/illustrationConfig/illustrationConfig'; + +import videosExploreService from 'services/api/videosExplore/videosExplore'; + +import { + decodeBufferPairs, + decodePathsVals, + iterFoldTree, +} from 'utils/encoder/streamEncoding'; + +import { ITraceVisualizerProps } from '../types'; + +import './VideosVisualizer.scss'; + +type VideoItem = { + blob_uri: string; + caption?: string; + context?: object; + format?: string; + fps?: number; + index?: number; + key?: string; + name?: string; + size?: number; + step?: number; +}; + +type BlobState = { + error?: string; + isLoading?: boolean; + url?: string; +}; + +function getVideoMimeType(format?: string): string { + switch ((format || '').toLowerCase()) { + case 'gif': + return 'image/gif'; + case 'm4v': + return 'video/mp4'; + case 'mov': + return 'video/quicktime'; + case 'webm': + return 'video/webm'; + case 'mp4': + default: + return 'video/mp4'; + } +} + +function formatBytes(bytes?: number): string | null { + if (!bytes) { + return null; + } + const units = ['B', 'KiB', 'MiB', 'GiB']; + let value = bytes; + let unitIndex = 0; + while (value >= 1024 && unitIndex < units.length - 1) { + value /= 1024; + unitIndex += 1; + } + return `${value.toFixed(unitIndex === 0 ? 0 : 1)} ${units[unitIndex]}`; +} + +function getBlobErrorStates(uris: string[], error: string) { + return uris.reduce>((states, uri) => { + states[uri] = { error, isLoading: false }; + return states; + }, {}); +} + +function VideosVisualizer({ + data, + isLoading, +}: ITraceVisualizerProps): React.FunctionComponentElement { + const videos: VideoItem[] = React.useMemo(() => data?.videos || [], [data]); + const [blobStates, setBlobStates] = React.useState>( + {}, + ); + const requestsRef = React.useRef void }>>({}); + const pendingVideosRef = React.useRef>({}); + const flushTimeoutRef = React.useRef(null); + const urlsRef = React.useRef>({}); + const formatsRef = React.useRef>({}); + + React.useEffect(() => { + videos.forEach((video) => { + formatsRef.current[video.blob_uri] = video.format; + }); + }, [videos]); + + React.useEffect(() => { + return () => { + if (flushTimeoutRef.current) { + window.clearTimeout(flushTimeoutRef.current); + } + Array.from(new Set(Object.values(requestsRef.current))).forEach( + (request) => request.abort(), + ); + Object.values(urlsRef.current).forEach((url) => URL.revokeObjectURL(url)); + requestsRef.current = {}; + pendingVideosRef.current = {}; + urlsRef.current = {}; + }; + }, []); + + const flushPendingVideoBlobRequests = React.useCallback(() => { + flushTimeoutRef.current = null; + + const pendingVideos = pendingVideosRef.current; + pendingVideosRef.current = {}; + const uris = Object.keys(pendingVideos).filter( + (uri) => !urlsRef.current[uri] && !requestsRef.current[uri], + ); + + if (!uris.length) { + return; + } + + const request = videosExploreService.getVideosByURIs(uris); + uris.forEach((uri) => { + requestsRef.current[uri] = request; + }); + request + .call() + .then(async (stream) => { + let bufferPairs = decodeBufferPairs(stream); + let decodedPairs = decodePathsVals(bufferPairs); + let objects = iterFoldTree(decodedPairs, 1); + const loadedUris = new Set(); + + for await (let [keys, val] of objects) { + const URI = keys[0] as string; + loadedUris.add(URI); + const blob = new Blob([val as ArrayBuffer], { + type: getVideoMimeType(formatsRef.current[URI]), + }); + const url = URL.createObjectURL(blob); + if (urlsRef.current[URI]) { + URL.revokeObjectURL(urlsRef.current[URI]); + } + urlsRef.current[URI] = url; + setBlobStates((state) => ({ + ...state, + [URI]: { isLoading: false, url }, + })); + } + + const missingUris = uris.filter((uri) => !loadedUris.has(uri)); + if (missingUris.length) { + setBlobStates((state) => ({ + ...state, + ...getBlobErrorStates(missingUris, 'Video blob not found'), + })); + } + }) + .catch((ex) => { + if (ex?.name !== 'AbortError') { + setBlobStates((state) => ({ + ...state, + ...getBlobErrorStates(uris, 'Could not load video'), + })); + } + }) + .finally(() => { + uris.forEach((uri) => { + if (requestsRef.current[uri] === request) { + delete requestsRef.current[uri]; + } + }); + }); + }, []); + + const requestVideoBlob = React.useCallback( + (video: VideoItem) => { + const uri = video.blob_uri; + if ( + !uri || + urlsRef.current[uri] || + requestsRef.current[uri] || + pendingVideosRef.current[uri] + ) { + return; + } + + formatsRef.current[uri] = video.format; + pendingVideosRef.current[uri] = video; + setBlobStates((state) => ({ + ...state, + [uri]: { ...state[uri], isLoading: true }, + })); + + if (!flushTimeoutRef.current) { + flushTimeoutRef.current = window.setTimeout( + flushPendingVideoBlobRequests, + 30, + ); + } + }, + [flushPendingVideoBlobRequests], + ); + + return ( + + +
+ {videos.length ? ( +
+ {videos.map((video) => ( + + ))} +
+ ) : ( + + )} +
+ + + ); +} + +function VideoCard({ + blobState, + onRequestBlob, + video, +}: { + blobState: BlobState; + onRequestBlob: (video: VideoItem) => void; + video: VideoItem; +}) { + const cardRef = React.useRef(null); + const videoRef = React.useRef(null); + const [isVisible, setIsVisible] = React.useState(false); + + React.useEffect(() => { + const card = cardRef.current; + if (!card) { + return; + } + if (!('IntersectionObserver' in window)) { + setIsVisible(true); + onRequestBlob(video); + return; + } + + const observer = new IntersectionObserver( + (entries) => { + const visible = entries[0]?.isIntersecting || false; + setIsVisible(visible); + if (visible) { + onRequestBlob(video); + } + }, + { rootMargin: '360px 0px', threshold: 0.1 }, + ); + observer.observe(card); + return () => observer.disconnect(); + }, [onRequestBlob, video]); + + React.useEffect(() => { + const player = videoRef.current; + if (!player) { + return; + } + if (isVisible && blobState.url) { + const playPromise = player.play(); + if (playPromise) { + playPromise.catch(() => {}); + } + } else { + player.pause(); + } + }, [blobState.url, isVisible]); + + const sizeText = formatBytes(video.size); + const isGif = (video.format || '').toLowerCase() === 'gif'; + const details = [ + typeof video.step === 'number' ? `step ${video.step}` : null, + typeof video.index === 'number' ? `index ${video.index}` : null, + video.fps ? `${video.fps} fps` : null, + sizeText, + ].filter(Boolean); + + return ( +
+
+ {isGif ? ( + {video.caption + ) : ( +
+
+ + {video.caption || video.name || 'Video'} + + {!!details.length && ( + + {details.join(' ยท ')} + + )} +
+
+ ); +} + +VideosVisualizer.displayName = 'VideosVisualizer'; + +export default React.memo(VideosVisualizer); diff --git a/aim/web/ui/src/pages/RunDetail/VideosVisualizer/index.ts b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/index.ts new file mode 100644 index 0000000000..8394ae0c24 --- /dev/null +++ b/aim/web/ui/src/pages/RunDetail/VideosVisualizer/index.ts @@ -0,0 +1,3 @@ +import VideosVisualizer from './VideosVisualizer'; + +export default VideosVisualizer; diff --git a/aim/web/ui/src/services/api/videosExplore/videosExplore.ts b/aim/web/ui/src/services/api/videosExplore/videosExplore.ts new file mode 100644 index 0000000000..c976255a44 --- /dev/null +++ b/aim/web/ui/src/services/api/videosExplore/videosExplore.ts @@ -0,0 +1,26 @@ +import { IApiRequest } from 'types/services/services'; + +import API from '../api'; + +const endpoints = { + GET_VIDEOS: 'runs/search/videos', + GET_VIDEOS_BY_URIS: 'runs/videos/get-batch', +}; + +function getVideosExploreData(params: {}): IApiRequest { + return API.getStream(endpoints.GET_VIDEOS, params); +} + +function getVideosByURIs(body: string[]): IApiRequest { + return API.getStream>(endpoints.GET_VIDEOS_BY_URIS, body, { + method: 'POST', + }); +} + +const videosExploreService = { + endpoints, + getVideosExploreData, + getVideosByURIs, +}; + +export default videosExploreService; diff --git a/aim/web/ui/src/services/models/runs/settings.ts b/aim/web/ui/src/services/models/runs/settings.ts index f98d7e2b8b..f473de666d 100644 --- a/aim/web/ui/src/services/models/runs/settings.ts +++ b/aim/web/ui/src/services/models/runs/settings.ts @@ -7,6 +7,7 @@ import { processTextsData, processPlotlyData, processAudiosData, + processVideosData, } from './util'; type InputItem = { @@ -113,6 +114,37 @@ const settings: Record = { }, }, }, + videos: { + dataProcessor: processVideosData, + sliders: { + record_range: { + defaultValue: [0, 50], + tooltip: 'Training step. Increments every time track() is called', + title: 'Steps', + sliderType: 'range', + infoPropertyName: 'step', + }, + index_range: { + defaultValue: [0, 50], + tooltip: 'Index in the list of videos passed to track() call', + title: 'Indices', + sliderType: 'range', + infoPropertyName: 'index', + }, + }, + inputs: { + record_density: { + defaultValue: 50, + title: 'Steps count', + tooltip: 'Number of steps to display', + }, + index_density: { + defaultValue: 5, + title: 'Indices count', + tooltip: 'Number of videos per step', + }, + }, + }, texts: { dataProcessor: processTextsData, sliders: { diff --git a/aim/web/ui/src/services/models/runs/util.ts b/aim/web/ui/src/services/models/runs/util.ts index cc85b64604..7b1f4bf13d 100644 --- a/aim/web/ui/src/services/models/runs/util.ts +++ b/aim/web/ui/src/services/models/runs/util.ts @@ -385,6 +385,50 @@ export function processAudiosData( }; } +export function processVideosData(data: Partial) { + const { + record_range_total, + iters, + values, + index_range_total, + context, + name, + } = data; + let videos: any[] = []; + + values?.forEach((stepData: IImageData[], stepIndex: number) => { + stepData.forEach((video: IImageData) => { + const step = iters?.[stepIndex]; + const videoKey = encode({ + name, + traceContext: context, + index: video.index, + step, + caption: video.caption, + }); + const seqKey = encode({ + name, + traceContext: context, + }); + videos.push({ + ...video, + name, + step, + context, + key: videoKey, + seqKey, + }); + }); + }); + + return { + videos: _.orderBy(videos, ['step', 'index'], ['desc', 'asc']), + record_range: [record_range_total?.[0], (record_range_total?.[1] || 0) - 1], + index_range: [index_range_total?.[0], (index_range_total?.[1] || 0) - 1], + processedDataType: VisualizationMenuTitles.videos, + }; +} + function groupData(data: IProcessedImageData[]): { key: string; config: { [key: string]: string }; diff --git a/docs/source/refs/sdk.rst b/docs/source/refs/sdk.rst index c217b3ef23..52b8f07984 100644 --- a/docs/source/refs/sdk.rst +++ b/docs/source/refs/sdk.rst @@ -45,6 +45,15 @@ aim.sdk.objects.audio .. autoclass:: Audio :members: +aim.sdk.objects.video +--------------------- + +.. automodule:: aim.sdk.objects.video +.. currentmodule:: aim.sdk.objects.video + +.. autoclass:: Video + :members: + aim.sdk.objects.text --------------------- @@ -112,6 +121,16 @@ aim.sdk.sequences.audio_sequence module :exclude-members: allowed_dtypes :members: +aim.sdk.sequences.video_sequence module +----------------------------- + +.. automodule:: aim.sdk.sequences.video_sequence +.. currentmodule:: aim.sdk.sequences.video_sequence + +.. autoclass:: Videos + :exclude-members: allowed_dtypes + :members: + aim.sdk.sequences.text_sequence module ----------------------------- diff --git a/pytest.ini b/pytest.ini index 8509e73b79..d8c428d5e4 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,4 @@ [pytest] addopts = --pdbcls=IPython.terminal.debugger:Pdb +markers = + system: slow end-to-end tests that may download or generate large external assets diff --git a/tests/api/test_run_images_api.py b/tests/api/test_run_images_api.py index c85518216b..7bc22521bc 100644 --- a/tests/api/test_run_images_api.py +++ b/tests/api/test_run_images_api.py @@ -1,11 +1,12 @@ import random +from aim.sdk.index_manager import RepoIndexManager from aim.sdk.run import Run from aim.storage.context import Context from aim.storage.treeutils import decode_tree from parameterized import parameterized from tests.base import ApiTestBase -from tests.utils import decode_encoded_tree_stream, generate_image_set +from tests.utils import decode_encoded_tree_stream, generate_image_set, generate_video_set class TestNoImagesRunQueryApi(ApiTestBase): @@ -355,6 +356,7 @@ def setUpClass(cls) -> None: super().setUpClass() # run1 -> context {'subset': 'train'} -> Image[] + # | -> Video[] # | -> integers # | -> floats # -> context {'subset': 'val'} -> floats @@ -363,23 +365,31 @@ def setUpClass(cls) -> None: # | -> floats # -> context {'subset': 'val'} -> floats - run1 = cls.create_run(system_tracking_interval=None) + run1 = cls.create_run(repo=cls.repo, system_tracking_interval=None) cls.run1_hash = run1.hash images = generate_image_set(img_count=2, caption_prefix='Image 0') + videos = generate_video_set(video_count=2, caption_prefix='Video 0') run1.track(images, name='image_lists', context={'subset': 'train'}) + run1.track(videos, name='video_lists', context={'subset': 'train'}) run1.track(random.random(), name='floats', context={'subset': 'train'}) run1.track(random.randint(100, 200), name='integers', context={'subset': 'train'}) run1.track(random.random(), name='floats', context={'subset': 'val'}) - run2 = cls.create_run(system_tracking_interval=None) + run2 = cls.create_run(repo=cls.repo, system_tracking_interval=None) run2.track(images[0], name='single_images', context={'subset': 'val'}) run2.track(random.random(), name='floats', context={'subset': 'train'}) run2.track(random.random(), name='floats', context={'subset': 'val'}) cls.run2_hash = run2.hash + run1.close() + run2.close() + index_manager = RepoIndexManager.get_index_manager(cls.repo) + index_manager.index(cls.run1_hash) + index_manager.index(cls.run2_hash) + cls.repo.container_pool.clear() def test_run_info_get_images_only_api(self): client = self.client - response = client.get(f'api/runs/{self.run1_hash}/info', params={'sequence': 'images'}) + response = client.get(f'/api/runs/{self.run1_hash}/info/', params={'sequence': 'images'}) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(1, len(response_data['traces'])) @@ -387,7 +397,7 @@ def test_run_info_get_images_only_api(self): self.assertDictEqual({'subset': 'train'}, response_data['traces']['images'][0]['context']) self.assertEqual('image_lists', response_data['traces']['images'][0]['name']) - response = client.get(f'api/runs/{self.run2_hash}/info', params={'sequence': 'images'}) + response = client.get(f'/api/runs/{self.run2_hash}/info/', params={'sequence': 'images'}) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(1, len(response_data['traces'])) @@ -398,19 +408,20 @@ def test_run_info_get_images_only_api(self): @parameterized.expand( [ ( - {'sequence': ('metric', 'images', 'audios', 'distributions', 'figures', 'texts', 'logs')}, - 7, + {'sequence': ('metric', 'images', 'videos', 'audios', 'distributions', 'figures', 'texts', 'logs')}, + 8, ), # explicit specification - (None, 8), # default + (None, 9), # default ] ) def test_run_info_get_all_sequences_api(self, qparams, trace_type_count): client = self.client - response = client.get(f'api/runs/{self.run1_hash}/info', params=qparams) + response = client.get(f'/api/runs/{self.run1_hash}/info/', params=qparams) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(trace_type_count, len(response_data['traces'])) self.assertIn('images', response_data['traces']) + self.assertIn('videos', response_data['traces']) self.assertIn('audios', response_data['traces']) self.assertIn('metric', response_data['traces']) self.assertIn('distributions', response_data['traces']) @@ -418,6 +429,8 @@ def test_run_info_get_all_sequences_api(self, qparams, trace_type_count): self.assertIn('texts', response_data['traces']) self.assertDictEqual({'subset': 'train'}, response_data['traces']['images'][0]['context']) self.assertEqual('image_lists', response_data['traces']['images'][0]['name']) + self.assertDictEqual({'subset': 'train'}, response_data['traces']['videos'][0]['context']) + self.assertEqual('video_lists', response_data['traces']['videos'][0]['name']) metrics_data = response_data['traces']['metric'] self.assertEqual(3, len(metrics_data)) self.assertEqual('floats', metrics_data[0]['name']) @@ -427,7 +440,7 @@ def test_run_info_get_all_sequences_api(self, qparams, trace_type_count): self.assertDictEqual({'subset': 'train'}, metrics_data[1]['context']) self.assertDictEqual({'subset': 'train'}, metrics_data[2]['context']) - response = client.get(f'api/runs/{self.run2_hash}/info', params={'sequence': ('images', 'metric')}) + response = client.get(f'/api/runs/{self.run2_hash}/info/', params={'sequence': ('images', 'metric')}) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(2, len(response_data['traces'])) @@ -444,14 +457,14 @@ def test_run_info_get_all_sequences_api(self, qparams, trace_type_count): def test_run_info_get_metrics_only_api(self): client = self.client - response = client.get(f'api/runs/{self.run1_hash}/info', params={'sequence': 'metric'}) + response = client.get(f'/api/runs/{self.run1_hash}/info/', params={'sequence': 'metric'}) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(1, len(response_data['traces'])) self.assertIn('metric', response_data['traces']) self.assertEqual(3, len(response_data['traces']['metric'])) - response = client.get(f'api/runs/{self.run2_hash}/info', params={'sequence': 'metric'}) + response = client.get(f'/api/runs/{self.run2_hash}/info/', params={'sequence': 'metric'}) self.assertEqual(200, response.status_code) response_data = response.json() self.assertEqual(1, len(response_data['traces'])) @@ -460,5 +473,5 @@ def test_run_info_get_metrics_only_api(self): def test_invalid_sequence_type(self): client = self.client - response = client.get(f'api/runs/{self.run1_hash}/info', params={'sequence': 'non-existing-sequence'}) + response = client.get(f'/api/runs/{self.run1_hash}/info/', params={'sequence': 'non-existing-sequence'}) self.assertEqual(400, response.status_code) diff --git a/tests/api/test_run_videos_api.py b/tests/api/test_run_videos_api.py new file mode 100644 index 0000000000..8fd1991cc9 --- /dev/null +++ b/tests/api/test_run_videos_api.py @@ -0,0 +1,146 @@ +import random + +from aim.sdk.index_manager import RepoIndexManager +from aim.sdk.run import Run +from aim.storage.context import Context +from aim.storage.treeutils import decode_tree +from parameterized import parameterized +from tests.base import ApiTestBase +from tests.utils import decode_encoded_tree_stream, generate_video_set + + +class TestNoVideosRunQueryApi(ApiTestBase): + def test_query_videos_api_empty_result(self): + client = self.client + + query = self.isolated_query_patch() + response = client.get('/api/runs/search/videos/', params={'q': query, 'report_progress': False}) + self.assertEqual(200, response.status_code) + self.assertEqual(b'', response.content) + + +class RunVideosTestBase(ApiTestBase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + run = cls.create_run(repo=cls.repo) + run['videos_per_step'] = 4 + for step in range(10): + videos = generate_video_set(video_count=4, caption_prefix=f'Video {step}') + run.track(videos, name='random_videos', step=step) + run.track(random.random(), name='random_values', step=step) + cls.run_hash = run.hash + run.close() + RepoIndexManager.get_index_manager(cls.repo).index(cls.run_hash) + cls.repo.container_pool.clear() + + +class TestRunVideosSearchApi(RunVideosTestBase): + def test_query_videos_api_defaults(self): + client = self.client + + query = self.isolated_query_patch() + response = client.get('/api/runs/search/videos/', params={'q': query, 'report_progress': False}) + self.assertEqual(200, response.status_code) + + decoded_response = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024))) + self.assertEqual(1, len(decoded_response)) + run_data = decoded_response[self.run_hash] + self.assertEqual([0, 10], run_data['ranges']['record_range_total']) + self.assertEqual([0, 10], run_data['ranges']['record_range_used']) + self.assertEqual([0, 4], run_data['ranges']['index_range_total']) + self.assertEqual([0, 4], run_data['ranges']['index_range_used']) + self.assertEqual(4, run_data['params']['videos_per_step']) + + trace_data = run_data['traces'][0] + self.assertEqual('random_videos', trace_data['name']) + self.assertEqual(10, len(trace_data['iters'])) + self.assertEqual(10, len(trace_data['values'])) + + video_list = trace_data['values'][2] + self.assertEqual(4, len(video_list)) + + video = video_list[3] + self.assertEqual('Video 2 3', video['caption']) + self.assertEqual(3, video['index']) + self.assertEqual('mp4', video['format']) + self.assertEqual(27, video['fps']) + self.assertEqual(len(b'fake mp4 bytes 3'), video['size']) + self.assertIn('blob_uri', video) + self.assertNotIn('data', video) + + def test_run_videos_batch_api(self): + client = self.client + + requested_traces = [ + {'name': 'random_videos', 'context': {}}, + ] + + response = client.post(f'/api/runs/{self.run_hash}/videos/get-batch/', json=requested_traces) + self.assertEqual(200, response.status_code) + + trace_data = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024))) + self.assertEqual('random_videos', trace_data['name']) + self.assertDictEqual({}, trace_data['context']) + self.assertEqual(10, len(trace_data['values'])) + self.assertEqual(list(range(10)), trace_data['iters']) + self.assertEqual('Video 2 3', trace_data['values'][2][3]['caption']) + self.assertIn('blob_uri', trace_data['values'][2][3]) + + +class RunVideosURIBulkLoadApi(RunVideosTestBase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + cls.video_blobs = {} + run = Run(run_hash=cls.run_hash, read_only=True) + empty_context = Context({}) + for step in range(3): + for idx in range(2): + video_view = run.series_run_trees[1].subtree((empty_context.idx, 'random_videos', 'val', step, idx)) + cls.video_blobs[video_view['caption']] = video_view['data'].load() + + @classmethod + def tearDownClass(cls) -> None: + cls.video_blobs.clear() + super().tearDownClass() + + def setUp(self) -> None: + super().setUp() + self.uri_map = {} + client = self.client + + response = client.get( + '/api/runs/search/videos/', + params={ + 'q': self.isolated_query_patch(), + 'record_range': '0:3', + 'index_range': '0:2', + 'record_density': 3, + 'index_density': 2, + 'report_progress': False, + }, + ) + decoded_response = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024))) + run_data = decoded_response[self.run_hash] + trace_data = run_data['traces'][0] + for video_list in trace_data['values']: + for video_data in video_list: + self.uri_map[video_data['blob_uri']] = video_data['caption'] + + def tearDown(self) -> None: + self.uri_map.clear() + super().tearDown() + + @parameterized.expand([(1,), (3,)]) + def test_videos_uri_bulk_load_api(self, uri_count): + uris = random.sample(list(self.uri_map.keys()), uri_count) + + client = self.client + response = client.post('/api/runs/videos/get-batch', json=uris) + self.assertEqual(200, response.status_code) + decoded_response = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024))) + self.assertEqual(uri_count, len(decoded_response)) + for uri, blob in decoded_response.items(): + expected_blob = self.video_blobs[self.uri_map[uri]] + self.assertEqual(expected_blob, blob) diff --git a/tests/sdk/test_video_construction.py b/tests/sdk/test_video_construction.py new file mode 100644 index 0000000000..619216f8cc --- /dev/null +++ b/tests/sdk/test_video_construction.py @@ -0,0 +1,64 @@ +import io +import os +import tempfile + +from aim.sdk import Video +from tests.base import TestBase + + +class TestVideoConstruction(TestBase): + def test_video_from_bytes(self): + video_bytes = b'fake mp4 bytes' + video = Video(data=video_bytes, format='mp4', fps=24, caption='sample') + + self.assertEqual('sample', video.caption) + self.assertEqual('mp4', video.format) + self.assertEqual(24, video.fps) + self.assertEqual(len(video_bytes), video.size) + self.assertEqual(video_bytes, video.get().read()) + self.assertDictEqual( + { + 'caption': 'sample', + 'format': 'mp4', + 'fps': 24, + 'size': len(video_bytes), + }, + video.json(), + ) + + def test_video_from_bytes_io(self): + video = Video(data=io.BytesIO(b'gif bytes'), format='gif') + + self.assertEqual('gif', video.format) + self.assertEqual(b'gif bytes', video.get().read()) + + def test_video_from_path_defers_blob_read_until_encode(self): + with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as file: + file.write(b'path video bytes') + video_path = file.name + + try: + video = Video(path=video_path, fps=30, caption='from path') + + self.assertEqual('mp4', video.format) + self.assertEqual(os.path.getsize(video_path), video.size) + self.assertEqual(os.path.abspath(video_path), video.storage['source_path']) + self.assertIsNone(video.storage.get('data')) + self.assertEqual(b'path video bytes', video.get().read()) + + name, _ = video._aim_encode() + + self.assertEqual('aim.video', name) + self.assertIsNone(video.storage.get('source_path')) + self.assertEqual(b'path video bytes', video.get().read()) + finally: + os.unlink(video_path) + + def test_video_requires_supported_format(self): + with self.assertRaises(ValueError): + Video(data=b'bytes', format='avi') + + def test_video_supports_common_mp4_container_extensions(self): + video = Video(data=b'bytes', format='m4v') + + self.assertEqual('m4v', video.format) diff --git a/tests/system/__init__.py b/tests/system/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/tests/system/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/system/test_video_upload_efficiency.py b/tests/system/test_video_upload_efficiency.py new file mode 100644 index 0000000000..542e4cdd9c --- /dev/null +++ b/tests/system/test_video_upload_efficiency.py @@ -0,0 +1,524 @@ +"""System coverage for first-class video tracking. + +The test is opt-in because it can download or generate many video samples. Set +``AIM_RUN_VIDEO_SYSTEM_TESTS=1`` to run it. If ``AIM_VIDEO_SYSTEM_DATASET_URL`` +points to a zip/tar archive or to a JSON manifest with ``videos`` entries, that +dataset is used. Otherwise ffmpeg is used to synthesize moving MP4 variants. +""" + +import gc +import hashlib +import json +import os +import shutil +import subprocess +import tarfile +import time +import urllib.parse +import urllib.request +import zipfile + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List + +import pytest + + +RUN_SYSTEM_TESTS_ENV = 'AIM_RUN_VIDEO_SYSTEM_TESTS' +DATASET_URL_ENV = 'AIM_VIDEO_SYSTEM_DATASET_URL' +CACHE_DIR_ENV = 'AIM_VIDEO_SYSTEM_CACHE_DIR' +VIDEO_COUNT_ENV = 'AIM_VIDEO_SYSTEM_VIDEO_COUNT' +MAX_RSS_BYTES_ENV = 'AIM_VIDEO_UPLOAD_MAX_RSS_BYTES' +MAX_STORAGE_MULTIPLIER_ENV = 'AIM_VIDEO_UPLOAD_MAX_STORAGE_MULTIPLIER' +MIN_UPLOAD_MIB_PER_SEC_ENV = 'AIM_VIDEO_UPLOAD_MIN_MIB_PER_SEC' +MAX_VIDEO_TRACK_P95_SECONDS_ENV = 'AIM_VIDEO_UPLOAD_MAX_TRACK_P95_SECONDS' + +DEFAULT_VIDEO_COUNT = 100 +MIB = 1024 * 1024 + + +@dataclass(frozen=True) +class VideoSample: + path: Path + width: int + height: int + fps: float + fmt: str + caption: str + + @property + def size_bytes(self) -> int: + return self.path.stat().st_size + + @property + def digest(self) -> str: + return _sha256_file(self.path) + + +@pytest.mark.system +def test_video_upload_many_files_is_efficient(tmp_path): + if os.environ.get(RUN_SYSTEM_TESTS_ENV) != '1': + pytest.skip(f'set {RUN_SYSTEM_TESTS_ENV}=1 to run the video upload system test') + + aim_module = pytest.importorskip('aim') + video_cls = getattr(aim_module, 'Video', None) + if video_cls is None: + pytest.skip('aim.Video is not implemented yet') + + from aim.sdk.configs import AIM_ENABLE_TRACKING_THREAD + from aim.sdk.index_manager import RepoIndexManager + from aim.sdk.repo import _get_tracking_queue + from aim.storage.treeutils import decode_tree + from aim.web.run import app + from fastapi.testclient import TestClient + from tests.utils import decode_encoded_tree_stream + + _assert_video_api_contract_is_registered() + + video_count = _int_env(VIDEO_COUNT_ENV, DEFAULT_VIDEO_COUNT) + assert video_count > 0 + samples = _prepare_video_samples(tmp_path, video_count) + _assert_dataset_shape(samples, video_count) + + previous_env_value = os.environ.get(AIM_ENABLE_TRACKING_THREAD) + previous_tracking_queue = aim_module.Repo.tracking_queue + os.environ[AIM_ENABLE_TRACKING_THREAD] = 'ON' + if aim_module.Repo.tracking_queue is None: + aim_module.Repo.tracking_queue = _get_tracking_queue() + + run = None + try: + repo = aim_module.Repo.default_repo() + run = aim_module.Run(repo=repo, system_tracking_interval=None) + assert run._tracker._non_blocking, 'video efficiency test must exercise async tracking' + run['testcase'] = __name__ + video_system_test_id = f'{int(time.time() * 1000)}-{os.getpid()}' + run['video_system_test_id'] = video_system_test_id + run_hash = run.hash + + source_total_bytes = sum(sample.size_bytes for sample in samples) + largest_video_bytes = max(sample.size_bytes for sample in samples) + repo_size_before = _directory_size(Path(repo.path)) + rss_before = _rss_bytes() + peak_rss = rss_before + started_at = time.perf_counter() + + blob_expectations: Dict[int, VideoSample] = {} + track_latencies = [] + for step, sample in enumerate(samples): + video = video_cls(path=str(sample.path), fps=sample.fps, format=sample.fmt, caption=sample.caption) + track_started_at = time.perf_counter() + run.track(video, name='stress_videos', step=step, context={'suite': 'video_system'}) + track_latencies.append(time.perf_counter() - track_started_at) + blob_expectations[step] = sample + del video + if step % 10 == 0: + gc.collect() + peak_rss = max(peak_rss, _rss_bytes()) + + logging_loop_elapsed = time.perf_counter() - started_at + if aim_module.Repo.tracking_queue is not None: + aim_module.Repo.tracking_queue._queue.join() + elapsed = time.perf_counter() - started_at + run.close() + RepoIndexManager.get_index_manager(repo).index(run_hash) + repo.container_pool.clear() + finally: + if previous_env_value is None: + os.environ.pop(AIM_ENABLE_TRACKING_THREAD, None) + else: + os.environ[AIM_ENABLE_TRACKING_THREAD] = previous_env_value + if previous_tracking_queue is not aim_module.Repo.tracking_queue: + aim_module.Repo.tracking_queue = previous_tracking_queue + + repo_size_after = _directory_size(Path(repo.path)) + gc.collect() + peak_rss = max(peak_rss, _rss_bytes()) + + max_video_track_p95 = _float_env(MAX_VIDEO_TRACK_P95_SECONDS_ENV, 0.25) + video_track_p95 = _percentile(track_latencies, 0.95) + assert video_track_p95 <= max_video_track_p95, ( + f'path-backed async video track() p95 was {video_track_p95 * 1000:.1f}ms; ' + f'limit is {max_video_track_p95 * 1000:.1f}ms' + ) + + max_rss_delta = _int_env(MAX_RSS_BYTES_ENV, max(512 * MIB, largest_video_bytes * 4)) + rss_delta = max(0, peak_rss - rss_before) + assert rss_delta <= max_rss_delta, ( + f'uploading {len(samples)} videos grew RSS by {rss_delta / MIB:.1f} MiB; ' + f'limit is {max_rss_delta / MIB:.1f} MiB' + ) + + storage_multiplier = _float_env(MAX_STORAGE_MULTIPLIER_ENV, 2.0) + max_storage_delta = int(source_total_bytes * storage_multiplier) + 128 * MIB + storage_delta = max(0, repo_size_after - repo_size_before) + assert storage_delta <= max_storage_delta, ( + f'video upload used {storage_delta / MIB:.1f} MiB for {source_total_bytes / MIB:.1f} MiB of source videos; ' + f'limit is {max_storage_delta / MIB:.1f} MiB' + ) + + min_upload_mib_per_sec = _float_env(MIN_UPLOAD_MIB_PER_SEC_ENV, 0.5) + upload_mib_per_sec = source_total_bytes / MIB / max(elapsed, 0.001) + assert upload_mib_per_sec >= min_upload_mib_per_sec, ( + f'video upload throughput was {upload_mib_per_sec:.2f} MiB/s; ' + f'limit is {min_upload_mib_per_sec:.2f} MiB/s' + ) + logging_loop_mib_per_sec = source_total_bytes / MIB / max(logging_loop_elapsed, 0.001) + assert logging_loop_mib_per_sec >= upload_mib_per_sec, 'async logging loop should not wait for full blob writes' + + try: + with TestClient(app) as client: + response = client.get( + '/api/runs/search/videos/', + params={ + 'q': f'run.video_system_test_id == "{video_system_test_id}"', + 'record_density': video_count, + 'index_density': 1, + 'report_progress': False, + }, + ) + assert response.status_code == 200, response.text + + decoded_response = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024))) + assert set(decoded_response.keys()) == {run_hash} + trace_data = decoded_response[run_hash]['traces'][0] + assert trace_data['name'] == 'stress_videos' + assert len(trace_data['values']) == video_count + + blob_uri_by_step = {} + for step, records in zip(trace_data['iters'], trace_data['values']): + assert len(records) == 1 + record = records[0] + assert 'blob_uri' in record + assert 'data' not in record + assert record.get('caption') == blob_expectations[step].caption + blob_uri_by_step[step] = record['blob_uri'] + + requested_steps = sorted({0, video_count // 2, video_count - 1}) + response = client.post('/api/runs/videos/get-batch', json=[blob_uri_by_step[step] for step in requested_steps]) + assert response.status_code == 200, response.text + + decoded_blobs = decode_tree( + decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024), concat_chunks=True) + ) + assert len(decoded_blobs) == len(requested_steps) + for step in requested_steps: + sample = blob_expectations[step] + blob = decoded_blobs[blob_uri_by_step[step]] + assert len(blob) == sample.size_bytes + assert hashlib.sha256(bytes(blob)).hexdigest() == sample.digest + finally: + repo.container_pool.clear() + + +def _assert_video_api_contract_is_registered() -> None: + from aim.sdk.sequence import Sequence + from aim.sdk.sequences.sequence_type_map import SEQUENCE_TYPE_MAP + + assert SEQUENCE_TYPE_MAP.get('aim.video') == 'videos' + assert SEQUENCE_TYPE_MAP.get('list(aim.video)') == 'videos' + assert 'videos' in Sequence.registry + + +def _prepare_video_samples(tmp_path: Path, video_count: int) -> List[VideoSample]: + dataset_url = os.environ.get(DATASET_URL_ENV) + if dataset_url: + dataset_dir = _download_and_unpack_dataset(dataset_url, tmp_path) + samples = _samples_from_dataset_dir(dataset_dir) + else: + samples = _generate_synthetic_video_dataset(tmp_path / 'synthetic-videos', video_count) + + return samples[:video_count] + + +def _download_and_unpack_dataset(dataset_url: str, tmp_path: Path) -> Path: + cache_root = Path(os.environ.get(CACHE_DIR_ENV, tmp_path / 'video-system-cache')).expanduser() + cache_root.mkdir(parents=True, exist_ok=True) + parsed_url = urllib.parse.urlparse(dataset_url) + archive_name = Path(parsed_url.path).name or 'video-dataset' + archive_path = cache_root / archive_name + dataset_dir = cache_root / archive_path.stem + + if not archive_path.exists(): + _download_file(dataset_url, archive_path) + + if archive_path.suffix == '.json': + dataset_dir.mkdir(parents=True, exist_ok=True) + _download_manifest_videos(archive_path, dataset_dir) + return dataset_dir + + if dataset_dir.exists() and any(dataset_dir.iterdir()): + return dataset_dir + + dataset_dir.mkdir(parents=True, exist_ok=True) + if zipfile.is_zipfile(archive_path): + with zipfile.ZipFile(archive_path) as archive: + _safe_extract_zip(archive, dataset_dir) + elif tarfile.is_tarfile(archive_path): + with tarfile.open(archive_path) as archive: + _safe_extract_tar(archive, dataset_dir) + else: + single_file = dataset_dir / archive_path.name + if not single_file.exists(): + shutil.copy2(archive_path, single_file) + return dataset_dir + + +def _download_manifest_videos(manifest_path: Path, dataset_dir: Path) -> None: + manifest = json.loads(manifest_path.read_text()) + for entry in manifest.get('videos', []): + video_url = entry['url'] + relative_path = Path(entry.get('path') or Path(urllib.parse.urlparse(video_url).path).name) + destination = dataset_dir / relative_path + destination.parent.mkdir(parents=True, exist_ok=True) + if not destination.exists(): + _download_file(video_url, destination) + shutil.copy2(manifest_path, dataset_dir / 'manifest.json') + + +def _download_file(url: str, destination: Path) -> None: + temporary_path = destination.with_suffix(destination.suffix + '.download') + with urllib.request.urlopen(url, timeout=60) as response, temporary_path.open('wb') as file: + shutil.copyfileobj(response, file, length=8 * MIB) + temporary_path.replace(destination) + + +def _safe_extract_zip(archive: zipfile.ZipFile, destination: Path) -> None: + destination_root = destination.resolve() + for member in archive.infolist(): + member_path = (destination / member.filename).resolve() + if destination_root not in (member_path, *member_path.parents): + raise ValueError(f'Unsafe path in zip dataset archive: {member.filename}') + archive.extractall(destination) + + +def _safe_extract_tar(archive: tarfile.TarFile, destination: Path) -> None: + destination_root = destination.resolve() + for member in archive.getmembers(): + member_path = (destination / member.name).resolve() + if destination_root not in (member_path, *member_path.parents): + raise ValueError(f'Unsafe path in tar dataset archive: {member.name}') + archive.extractall(destination) + + +def _samples_from_dataset_dir(dataset_dir: Path) -> List[VideoSample]: + metadata_by_name = _load_manifest_metadata(dataset_dir) + video_paths = sorted( + path + for path in dataset_dir.rglob('*') + if path.suffix.lower() in {'.mp4', '.m4v', '.mov', '.webm', '.gif'} + ) + samples = [] + for idx, path in enumerate(video_paths): + metadata = metadata_by_name.get(path.name) or metadata_by_name.get(path.relative_to(dataset_dir).as_posix()) + if metadata: + width = int(metadata['width']) + height = int(metadata['height']) + fps = float(metadata['fps']) + else: + width, height, fps = _probe_video(path) + samples.append( + VideoSample( + path=path, + width=width, + height=height, + fps=fps, + fmt=path.suffix.lower().lstrip('.'), + caption=f'video-system-{idx}-{width}x{height}-{fps:g}fps', + ) + ) + return samples + + +def _load_manifest_metadata(dataset_dir: Path) -> Dict[str, dict]: + manifest_path = dataset_dir / 'manifest.json' + if not manifest_path.exists(): + return {} + manifest = json.loads(manifest_path.read_text()) + metadata = {} + for entry in manifest.get('videos', []): + key = entry.get('path') or Path(urllib.parse.urlparse(entry.get('url', '')).path).name + if key and {'width', 'height', 'fps'} <= set(entry.keys()): + metadata[key] = entry + metadata[Path(key).name] = entry + return metadata + + +def _generate_synthetic_video_dataset(dataset_dir: Path, video_count: int) -> List[VideoSample]: + ffmpeg = shutil.which('ffmpeg') + if not ffmpeg: + pytest.skip(f'{DATASET_URL_ENV} is not set and ffmpeg is not available to generate a video dataset') + + dataset_dir.mkdir(parents=True, exist_ok=True) + source_path = dataset_dir / 'synthetic_source_1920x1080_60fps.mp4' + if not source_path.exists(): + _run_ffmpeg_source(ffmpeg, source_path) + variants = [ + (320, 200, 12), + (426, 240, 15), + (640, 360, 24), + (1280, 720, 30), + (1920, 1080, 60), + ] + samples = [] + for idx in range(video_count): + variant_idx = idx % len(variants) + width, height, fps = variants[variant_idx] + output_path = dataset_dir / f'synthetic_variant_{variant_idx}_{width}x{height}_{fps}fps.mp4' + if not output_path.exists(): + _run_ffmpeg_variant(ffmpeg, source_path, output_path, width, height, fps) + samples.append( + VideoSample( + path=output_path, + width=width, + height=height, + fps=float(fps), + fmt='mp4', + caption=f'video-system-{idx}-{width}x{height}-{fps}fps', + ) + ) + return samples + + +def _run_ffmpeg_source(ffmpeg: str, output_path: Path) -> None: + base_cmd = [ + ffmpeg, + '-hide_banner', + '-loglevel', + 'error', + '-y', + '-f', + 'lavfi', + '-i', + 'testsrc2=size=1920x1080:rate=60:duration=1.0', + '-pix_fmt', + 'yuv420p', + '-an', + '-movflags', + '+faststart', + ] + _run_ffmpeg_with_fallbacks( + base_cmd, + output_path, + (['-c:v', 'libx264', '-preset', 'ultrafast', '-crf', '30'], ['-c:v', 'mpeg4', '-q:v', '6']), + ) + + +def _run_ffmpeg_variant(ffmpeg: str, source_path: Path, output_path: Path, width: int, height: int, fps: int) -> None: + base_cmd = [ + ffmpeg, + '-hide_banner', + '-loglevel', + 'error', + '-y', + '-i', + str(source_path), + '-vf', + f'scale={width}:{height}', + '-r', + str(fps), + '-pix_fmt', + 'yuv420p', + '-an', + '-movflags', + '+faststart', + ] + _run_ffmpeg_with_fallbacks( + base_cmd, + output_path, + (['-c:v', 'libx264', '-preset', 'ultrafast', '-crf', '30'], ['-c:v', 'mpeg4', '-q:v', '6']), + ) + + +def _run_ffmpeg_with_fallbacks(base_cmd: List[str], output_path: Path, codec_arg_sets: tuple) -> None: + result = None + for codec_args in codec_arg_sets: + result = subprocess.run(base_cmd + codec_args + [str(output_path)], capture_output=True, text=True) + if result.returncode == 0: + return + pytest.fail(f'ffmpeg failed to generate {output_path}: {result.stderr}') + + +def _probe_video(path: Path) -> tuple: + ffprobe = shutil.which('ffprobe') + if not ffprobe: + pytest.skip('downloaded dataset has no manifest metadata and ffprobe is not available') + result = subprocess.run( + [ + ffprobe, + '-v', + 'error', + '-select_streams', + 'v:0', + '-show_entries', + 'stream=width,height,r_frame_rate', + '-of', + 'json', + str(path), + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + pytest.fail(f'ffprobe failed for {path}: {result.stderr}') + stream = json.loads(result.stdout)['streams'][0] + fps = _parse_fps(stream['r_frame_rate']) + return int(stream['width']), int(stream['height']), fps + + +def _assert_dataset_shape(samples: List[VideoSample], expected_count: int) -> None: + assert len(samples) == expected_count + resolutions = {(sample.width, sample.height) for sample in samples} + fps_values = {round(sample.fps, 2) for sample in samples} + assert len(resolutions) >= _int_env('AIM_VIDEO_SYSTEM_MIN_RESOLUTION_BUCKETS', 5) + assert len(fps_values) >= _int_env('AIM_VIDEO_SYSTEM_MIN_FPS_BUCKETS', 3) + assert any(sample.width <= 320 and sample.height <= 240 for sample in samples) + assert any(sample.width >= 1920 and sample.height >= 1080 for sample in samples) + + +def _directory_size(path: Path) -> int: + total = 0 + if not path.exists(): + return total + for file_path in path.rglob('*'): + if file_path.is_file(): + total += file_path.stat().st_size + return total + + +def _rss_bytes() -> int: + try: + import psutil + except ImportError: + return 0 + return psutil.Process(os.getpid()).memory_info().rss + + +def _sha256_file(path: Path) -> str: + hasher = hashlib.sha256() + with path.open('rb') as file: + for chunk in iter(lambda: file.read(8 * MIB), b''): + hasher.update(chunk) + return hasher.hexdigest() + + +def _parse_fps(raw: str) -> float: + numerator, _, denominator = raw.partition('/') + if denominator: + return float(numerator) / float(denominator) + return float(numerator) + + +def _int_env(name: str, default: int) -> int: + return int(os.environ.get(name, default)) + + +def _float_env(name: str, default: float) -> float: + return float(os.environ.get(name, default)) + + +def _percentile(values: List[float], pct: float) -> float: + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int(round((len(ordered) - 1) * pct)))) + return ordered[index] diff --git a/tests/utils.py b/tests/utils.py index b10a7a7954..cfb7e20398 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -7,6 +7,7 @@ import numpy from aim.sdk.objects.image import Image as AimImage +from aim.sdk.objects.video import Video as AimVideo from aim.sdk.repo import Repo from aim.sdk.run import Run from aim.storage.structured.sql_engine.models import Base as StructuredBase @@ -63,6 +64,18 @@ def generate_image_set(img_count, caption_prefix='Image', img_size=(16, 16)): ] +def generate_video_set(video_count, caption_prefix='Video'): + return [ + AimVideo( + data=f'fake mp4 bytes {idx}'.encode(), + format='mp4', + fps=24 + idx, + caption=f'{caption_prefix} {idx}', + ) + for idx in range(video_count) + ] + + def truncate_structured_db(db): session = db.get_session() meta = StructuredBase.metadata