# Copyright 2022 AI Singapore
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Writes the output image/video to file.
"""
import datetime
from pathlib import Path
from typing import Any, Dict, Optional
import cv2
import numpy as np
from peekingduck.pipeline.nodes.abstract_node import AbstractNode
# role of this node is to be able to take in multiple frames, stitch them
# together and output them.
# to do: need to have 'live' kind of data when there is no filename
# to do: it will be good to have the accepted file format as a configuration
# to do: somewhere so that input and output can use this config for media related issues
[docs]class Node(AbstractNode):
"""Outputs the processed image or video to a file. A timestamp is appended to the
end of the file name.
Inputs:
|img_data|
|filename_data|
|saved_video_fps_data|
|pipeline_end_data|
Outputs:
|none_output_data|
Configs:
output_dir (:obj:`str`): **default = "PeekingDuck/data/output"**. |br|
Output directory for files to be written locally.
"""
def __init__(self, config: Dict[str, Any] = None, **kwargs: Any) -> None:
super().__init__(config, node_path=__name__, **kwargs)
self.output_dir = Path(self.output_dir) # type: ignore
self._file_name: Optional[str] = None
self._file_path_with_timestamp: Optional[str] = None
self._image_type: Optional[str] = None
self.writer = None
self._prepare_directory(self.output_dir)
self._fourcc = cv2.VideoWriter_fourcc(*"mp4v")
self.logger.info(f"Output directory used is: {self.output_dir}")
def run(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
"""Writes media information to filepath."""
# reset and terminate when there are no more data
if inputs["pipeline_end"]:
if self.writer: # images automatically releases writer
self.writer.release()
return {}
if not self._file_name:
self._prepare_writer(
inputs["filename"], inputs["img"], inputs["saved_video_fps"]
)
if inputs["filename"] != self._file_name:
self._prepare_writer(
inputs["filename"], inputs["img"], inputs["saved_video_fps"]
)
self._write(inputs["img"])
return {}
def _get_config_types(self) -> Dict[str, Any]:
"""Returns dictionary mapping the node's config keys to respective types."""
return {"output_dir": str}
def _write(self, img: np.ndarray) -> None:
if self._image_type == "image":
cv2.imwrite(self._file_path_with_timestamp, img)
else:
self.writer.write(img)
def _prepare_writer(
self, filename: str, img: np.ndarray, saved_video_fps: int
) -> None:
self._file_path_with_timestamp = self._append_datetime_filename(filename)
if filename.split(".")[-1] in ["jpg", "jpeg", "png"]:
self._image_type = "image"
else:
self._image_type = "video"
resolution = img.shape[1], img.shape[0]
self.writer = cv2.VideoWriter(
self._file_path_with_timestamp,
self._fourcc,
saved_video_fps,
resolution,
)
@staticmethod
def _prepare_directory(output_dir: Path) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
def _append_datetime_filename(self, filename: str) -> str:
self._file_name = filename
current_time = datetime.datetime.now()
# output as 'YYYYMMDD_hhmmss'
time_str = current_time.strftime("%y%m%d_%H%M%S")
# append timestamp to filename before extension Format: filename_timestamp.extension
filename_with_timestamp = f"_{time_str}.".join(filename.split(".")[-2:])
file_path_with_timestamp = self.output_dir / filename_with_timestamp
return str(file_path_with_timestamp)