platypush/platypush/plugins/camera/model/writer/cv.py

46 lines
1.2 KiB
Python

import numpy as np
from PIL.Image import Image as ImageType
from platypush.plugins.camera.model.writer import FileVideoWriter
class CvFileWriter(FileVideoWriter):
"""
Write camera frames to a file using OpenCV.
"""
def __init__(self, *args, **kwargs):
import cv2
super(CvFileWriter, self).__init__(*args, **kwargs)
video_type = cv2.VideoWriter_fourcc(*(self.camera.info.output_format or 'xvid').upper())
resolution = (
int(self.camera.info.resolution[0] * (self.camera.info.scale_x or 1.)),
int(self.camera.info.resolution[1] * (self.camera.info.scale_y or 1.)),
)
self.writer = cv2.VideoWriter(self.output_file, video_type, self.camera.info.fps, resolution, False)
def write(self, img):
if not self.writer:
return
try:
if isinstance(img, ImageType):
# noinspection PyTypeChecker
img = np.array(img)
except (ValueError, TypeError):
pass
self.writer.write(img)
def close(self):
if not self.writer:
return
self.writer.release()
self.writer = None
super().close()
# vim:sw=4:ts=4:et: