diff --git a/vot/report/__init__.py b/vot/report/__init__.py index fb5a59f..28bc76f 100644 --- a/vot/report/__init__.py +++ b/vot/report/__init__.py @@ -125,43 +125,34 @@ def save(self, output: str, fmt: str): import tempfile import shutil import os - import cv2 + from .video import VideoWriterScikitH264, VideoWriterOpenCV supported_mappings = { - "mp4": "mp4v", - "avi": "XVID" + "mp4": VideoWriterScikitH264, + "avi": VideoWriterOpenCV } if not fmt in supported_mappings: raise ValueError("Unsupported video format: {}".format(fmt)) - fourcc = cv2.VideoWriter_fourcc(*supported_mappings[fmt]) - - frame = self.render(0) - height, width, _ = frame.shape - if not isinstance(output, str): - fd, tempname = tempfile.mkstemp() + fd, tempname = tempfile.mkstemp(prefix="video_", suffix=".{}".format(fmt)) os.close(fd) else: tempname = output - writer = cv2.VideoWriter(tempname, fourcc, self._fps, (height, width)) - print(frame.dtype, frame.shape) - writer.write(frame) + writer = supported_mappings[fmt](tempname, self._fps) - for i in range(1, len(self._frames)): - frame = self.render(i) + for i in range(0, len(self._frames)): + writer(self.render(i)) - writer.write(frame) - - writer.release() + writer.close() if tempname == output: return shutil.copyfileobj(open(tempname, 'rb'), output) - #os.remove(tempname) + os.remove(tempname) def __len__(self): return len(self._frames) @@ -205,8 +196,8 @@ def draw(self, key, data): class ObjectVideo(Video): - def __init__(self, identifier: str, frames: FrameList, trait=None): - super().__init__(identifier, frames, trait) + def __init__(self, identifier: str, frames: FrameList, fps=10, trait=None): + super().__init__(identifier, frames, fps=fps, trait=trait) self._regions = {} def draw(self, frame, key, data): @@ -753,30 +744,32 @@ def update(): report_storage = workspace.storage.substorage("reports").substorage(name) - def only_plots(reports, format: str, storage: "Storage"): + def only_plots(reports, storage: "Storage"): """Filter out all non-plot items from the report and save them to storage. Args: reports: The reports to filter. - format: The format to save the plots in. """ for key, section in reports.items(): for item in section: if isinstance(item, Plot): logger.debug("Saving plot %s", item.identifier) - plot_name = key + "_" + item.identifier + '.%s' % format.lower() - with storage.write(plot_name, binary=True) as out: - item.save(out, format.upper()) - + with storage.write(key + "_" + item.identifier + '.pdf', binary=True) as out: + item.save(out, "PDF") + with storage.write(key + "_" + item.identifier + '.png', binary=True) as out: + item.save(out, "PNG") + if isinstance(item, Video): + logger.debug("Saving video %s", item.identifier) + with storage.write(key + "_" + item.identifier + '.avi', binary=True) as out: + item.save(out, "avi") + if format == "html": from .html import generate_html_document generate_html_document(trackers, workspace.dataset, reports, report_storage) elif format == "latex": from .latex import generate_latex_document generate_latex_document(trackers, workspace.dataset, reports, report_storage) - elif format == "pdf_plots": - only_plots(reports, "pdf", report_storage) - elif format == "png_plots": - only_plots(reports, "png", report_storage) + elif format == "plots": + only_plots(reports, report_storage) else: raise ValueError("Unknown report format %s" % format) \ No newline at end of file diff --git a/vot/report/html.py b/vot/report/html.py index f32dd4f..c54f5fe 100644 --- a/vot/report/html.py +++ b/vot/report/html.py @@ -71,9 +71,10 @@ def generate_html_document(trackers: List[Tracker], sequences: List[Sequence], r def insert_video(data: Video): """Insert a video into the document.""" - name = data.identifier + ".avi" + name = data.identifier + ".mp4" - data.save(storage.write(name), "avi") + with storage.write(name, binary=True) as handle: + data.save(handle, "mp4") with video(src=name, controls=True, preload="auto", autoplay=False, loop=False, width="100%", height="100%"): raw("Your browser does not support the video tag.") diff --git a/vot/report/video.py b/vot/report/video.py new file mode 100644 index 0000000..e96481c --- /dev/null +++ b/vot/report/video.py @@ -0,0 +1,122 @@ + + +from typing import List + +from attributee import Boolean + +from vot.dataset import Sequence +from vot.tracker import Tracker +from vot.experiment.multirun import MultiRunExperiment, Experiment +from vot.report import ObjectVideo, SeparableReport + +class VideoWriter: + + def __init__(self, filename: str, fps: int = 30): + self._filename = filename + self._fps = fps + + def __call__(self, frame): + raise NotImplementedError() + + def close(self): + raise NotImplementedError() + +class VideoWriterScikitH264(VideoWriter): + + def _handle(self): + try: + import skvideo.io + except ImportError: + raise ImportError("The scikit-video package is required for video export.") + if not hasattr(self, "_writer"): + skvideo.io.vwrite(self._filename, []) + self._writer = skvideo.io.FFmpegWriter(self._filename, inputdict={'-r': str(self._fps), '-vcodec': 'libx264'}) + return self._writer + + def __call__(self, frame): + self._handle().writeFrame(frame) + + def close(self): + if hasattr(self, "_writer"): + self._writer.close() + self._writer = None + +class VideoWriterOpenCV(VideoWriter): + + + def __init__(self, filename: str, fps: int = 30, codec: str = "mp4v"): + super().__init__(filename, fps) + self._codec = codec + + def __call__(self, frame): + try: + import cv2 + except ImportError: + raise ImportError("The OpenCV package is required for video export.") + if not hasattr(self, "_writer"): + self._height, self._width = frame.shape[:2] + self._writer = cv2.VideoWriter(self._filename, cv2.VideoWriter_fourcc(*self._codec.lower()), self._fps, (self._width, self._height)) + self._writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + + def close(self): + if hasattr(self, "_writer"): + self._writer.release() + self._writer = None + +class PreviewVideos(SeparableReport): + """A report that generates video previews for the tracker results.""" + + groundtruth = Boolean(default=False, description="If set, the groundtruth is shown with the tracker output.") + separate = Boolean(default=False, description="If set, each tracker is shown in a separate video.") + + async def perexperiment(self, experiment: Experiment, trackers: List[Tracker], sequences: List[Sequence]): + + videos = [] + + for sequence in sequences: + + if self.separate: + + for tracker in trackers: + + video = ObjectVideo(sequence.identifier + "_" + tracker.identifier, sequence) + + if self.groundtruth: + for frame in range(len(sequence)): + video(frame, "_", sequence.groundtruth(frame)) + + for obj in sequence.objects(): + trajectories = experiment.gather(tracker, sequence, objects=[obj]) + + if len(trajectories) == 0: + continue + + for frame in range(len(sequence)): + video(frame, obj, trajectories[0].region(frame)) + + videos.append(video) + else: + + video = ObjectVideo(sequence.identifier, sequence) + + if self.groundtruth: + for frame in range(len(sequence)): + video(frame, "_", sequence.groundtruth(frame)) + + for tracker in trackers: + + for obj in sequence.objects(): + trajectories = experiment.gather(tracker, sequence, objects=[obj]) + + if len(trajectories) == 0: + continue + + for frame in range(len(sequence)): + video(frame, obj + "_" + tracker.identifier, trajectories[0].region(frame)) + + videos.append(video) + + return videos + + def compatible(self, experiment): + return isinstance(experiment, MultiRunExperiment) \ No newline at end of file diff --git a/vot/report/videos.py b/vot/report/videos.py deleted file mode 100644 index e90e51a..0000000 --- a/vot/report/videos.py +++ /dev/null @@ -1,46 +0,0 @@ - - -from typing import List - -from attributee import Boolean - -from vot.dataset import Sequence -from vot.tracker import Tracker -from vot.experiment.multirun import MultiRunExperiment, Experiment -from vot.report import ObjectVideo, SeparableReport - - -class PreviewVideos(SeparableReport): - """A report that generates video previews for the tracker results.""" - - groundtruth = Boolean(default=False, description="If set, the groundtruth is shown with the tracker output.") - - async def perexperiment(self, experiment: Experiment, trackers: List[Tracker], sequences: List[Sequence]): - - videos = [] - - for sequence in sequences: - - for tracker in trackers: - - video = ObjectVideo(sequence.identifier + "_" + tracker.identifier, sequence) - - if self.groundtruth: - for frame in range(len(sequence)): - video(frame, "_", sequence.groundtruth(frame)) - - for obj in sequence.objects(): - trajectories = experiment.gather(tracker, sequence, objects=[obj]) - - if len(trajectories) == 0: - continue - - for frame in range(len(sequence)): - video(frame, obj, trajectories[0].region(frame)) - - videos.append(video) - - return videos - - def compatible(self, experiment): - return isinstance(experiment, MultiRunExperiment) \ No newline at end of file