"""Streamer for reading input."""
# Copyright (C) 2021-2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#
import abc
import multiprocessing
import os
import queue
import sys
from enum import Enum
from typing import Iterator, Union
import cv2
import numpy as np
[docs]
class OpenError(Exception):
"""Exception for error opening reader."""
def __init__(self, message: str) -> None:
super().__init__(message)
self.message = message
class MediaType(Enum):
"""This Enum represents the types of input."""
IMAGE = 1
DIR = 2
VIDEO = 3
CAMERA = 4
[docs]
class BaseStreamer(metaclass=abc.ABCMeta):
"""Base Streamer interface to implement Image, Video and Camera streamers."""
@abc.abstractmethod
def __iter__(self) -> Iterator[np.ndarray]:
"""Iterate through the streamer object that is a Python Generator object.
Returns:
np.ndarray: Yield the image or video frame.
"""
raise NotImplementedError
[docs]
@abc.abstractmethod
def get_type(self) -> MediaType:
"""Get type of streamer.
Returns:
MediaType: type of streamer.
"""
raise NotImplementedError
[docs]
def fps(self):
"""Returns a frequency of getting images from source."""
raise NotImplementedError
def _process_run(streamer: BaseStreamer, buffer: multiprocessing.Queue) -> None:
"""Private function that is run by the thread.
Waits for the buffer to gain space for timeout seconds while it is full.
If no space was available within this time the function will exit
streamer (BaseStreamer): The streamer to retrieve frames from
buffer (multiprocessing.Queue): The buffer to place the retrieved frames in
"""
for frame in streamer:
buffer.put(frame)
[docs]
class ThreadedStreamer(BaseStreamer):
"""Runs a BaseStreamer on a separate thread.
streamer (BaseStreamer): The streamer to run on a thread
buffer_size (int): Number of frame to buffer internally. Defaults to 2.
Example:
>>> streamer = VideoStreamer(path="../demo.mp4")
>>> threaded_streamer = ThreadedStreamer(streamer)
>>> for frame in threaded_streamer:
... pass
"""
def __init__(self, streamer: BaseStreamer, buffer_size: int = 2) -> None:
self.buffer_size = buffer_size
self.streamer = streamer
def __iter__(self) -> Iterator[np.ndarray]:
"""Get frames from streamer and yield them.
Yields:
Iterator[np.ndarray]: Yield the image or video frame.
"""
buffer: multiprocessing.Queue = multiprocessing.Queue(maxsize=self.buffer_size)
process = multiprocessing.Process(target=_process_run, args=(self.streamer, buffer))
# Make thread a daemon so that it will exit when the main program exits as well
process.daemon = True
process.start()
try:
while process.is_alive() or not buffer.empty():
try:
yield buffer.get(timeout=0.1)
except queue.Empty:
pass
except GeneratorExit:
process.terminate()
finally:
process.join(timeout=0.1)
# The kill() function is only available in Python 3.7.
# Skip it if running an older Python version.
if sys.version_info >= (3, 7) and process.exitcode is None:
process.kill()
[docs]
def get_type(self) -> MediaType:
"""Get type of internal streamer.
Returns:
MediaType: type of internal streamer.
"""
return self.streamer.get_type()
[docs]
class VideoStreamer(BaseStreamer):
"""Video Streamer.
Args:
path: Path to the video file.
Example:
>>> streamer = VideoStreamer(path="../demo.mp4")
... for frame in streamer:
... pass
"""
def __init__(self, input_path: str, loop: bool = False) -> None:
self.media_type = MediaType.VIDEO
self.loop = loop
self.cap = cv2.VideoCapture()
status = self.cap.open(input_path)
if not status:
raise InvalidInput(f"Can't open the video from {input_path}")
def __iter__(self) -> Iterator[np.ndarray]:
"""Iterates over frames of the video.
If self.loop is set to True, the video will loop infinitely.
"""
while True:
status, image = self.cap.read()
if status:
yield cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
else:
if self.loop:
self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
else:
break
[docs]
def fps(self):
"""Returns a frequency of getting images from source."""
return self.cap.get(cv2.CAP_PROP_FPS)
[docs]
def get_type(self) -> MediaType:
"""Returns the type of media."""
return MediaType.VIDEO
[docs]
class CameraStreamer(BaseStreamer):
"""Stream video frames from camera.
Args:
camera_device (int): Camera device index e.g, 0, 1
Example:
>>> streamer = CameraStreamer(camera_device=0)
... for frame in streamer:
... cv2.imshow("Window", frame)
... if ord("q") == cv2.waitKey(1):
... break
"""
def __init__(self, camera_device: int = 0) -> None:
self.media_type = MediaType.CAMERA
try:
self.stream = cv2.VideoCapture(int(camera_device))
except ValueError as error:
raise InvalidInput(f"Can't find the camera {camera_device}") from error
def __iter__(self) -> Iterator[np.ndarray]:
"""Read video and yield the frame.
Args:
stream: Video stream captured via OpenCV's VideoCapture
Returns:
Individual frame
"""
while True:
frame_available, frame = self.stream.read()
if not frame_available:
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
yield frame
self.stream.release()
[docs]
def get_type(self) -> MediaType:
"""Returns the type of media."""
return MediaType.CAMERA
[docs]
class ImageStreamer(BaseStreamer):
"""Stream from image file.
Args:
input_path (str): Path to an image.
loop (bool): Whether to loop through the image or not. Defaults to False.
Example:
>>> streamer = ImageStreamer(path="../images")
... for frame in streamer:
... cv2.imshow("Window", frame)
... cv2.waitKey(0)
"""
def __init__(self, input_path: str, loop: bool = False) -> None:
self.loop = loop
self.media_type = MediaType.IMAGE
if not os.path.isfile(input_path):
raise InvalidInput(f"Can't find the image by {input_path}")
self.image = cv2.imread(input_path, cv2.IMREAD_COLOR)
if self.image is None:
raise OpenError(f"Can't open the image from {input_path}")
self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
def __iter__(self) -> Iterator[np.ndarray]:
"""If loop is True, yield the image again and again."""
if not self.loop:
yield self.image
else:
while True:
yield self.image
[docs]
def get_type(self) -> MediaType:
"""Returns the type of the streamer."""
return MediaType.IMAGE
[docs]
class DirStreamer(BaseStreamer):
"""Stream from directory of images.
Args:
path: Path to directory.
Example:
>>> streamer = DirStreamer(path="../images")
... for frame in streamer:
... cv2.imshow("Window", frame)
... cv2.waitKey(0)
"""
def __init__(self, input_path: str, loop: bool = False) -> None:
self.loop = loop
self.media_type = MediaType.DIR
self.dir = input_path
if not os.path.isdir(self.dir):
raise InvalidInput(f"Can't find the dir by {input_path}")
self.names = sorted(os.listdir(self.dir))
if not self.names:
raise OpenError(f"The dir {input_path} is empty")
self.file_id = 0
for name in self.names:
filename = os.path.join(self.dir, name)
image = cv2.imread(str(filename), cv2.IMREAD_COLOR)
if image is not None:
return
raise OpenError(f"Can't read the first image from {input_path}")
def __iter__(self) -> Iterator[np.ndarray]:
"""Iterates over the images in a directory.
If self.loop is True, it reiterates again from the first image in the directory.
"""
while self.file_id < len(self.names):
filename = os.path.join(self.dir, self.names[self.file_id])
image = cv2.imread(str(filename), cv2.IMREAD_COLOR)
if self.file_id < len(self.names) - 1:
self.file_id = self.file_id + 1
else:
self.file_id = self.file_id + 1 if not self.loop else 0
if image is not None:
yield cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
[docs]
def get_type(self) -> MediaType:
"""Returns the type of the streamer."""
return MediaType.DIR
[docs]
def get_streamer(
input_stream: Union[int, str] = 0,
loop: bool = False,
threaded: bool = False,
) -> BaseStreamer:
"""Get streamer object based on the file path or camera device index provided.
Args:
input_stream (Union[int, str]): Path to file or directory or index for camera.
loop (bool): Enable reading the input in a loop. Defaults to False.
threaded (bool): Run streaming on a separate thread. Threaded streaming option. Defaults to False.
Returns:
BaseStreamer: Streamer object.
"""
# errors: Dict = {InvalidInput: [], OpenError: []}
errors = []
streamer: BaseStreamer
for reader in (ImageStreamer, DirStreamer, VideoStreamer):
try:
streamer = reader(input_stream, loop) # type: ignore
if threaded:
streamer = ThreadedStreamer(streamer)
return streamer
except (InvalidInput, OpenError) as error:
errors.append(error)
try:
streamer = CameraStreamer(input_stream) # type: ignore
if threaded:
streamer = ThreadedStreamer(streamer)
return streamer
except (InvalidInput, OpenError) as error:
errors.append(error)
if errors:
raise Exception(errors)
sys.exit(1)
