""" Common function for camera based method """ from fractions import Fraction import json import logging import os import time import av import cv2 import numpy as np from settings.config import settings logger = logging.getLogger(__name__) class VideoAnalyser(object): """摄像头/视频数据分析的基类, 实现逐帧分析 Attributes: t_start_save_video (float): 开始保存视频的时间，当使用av保存视频时需要此参数计算pts out_stream: 使用opencv保存数据时使用 container：使用av保存视频时使用 stream： 使用av保存时使用 """ def __init__(self, camera_id=0, input_video=None): if not input_video: # For webcam input: self.camera_id = camera_id self.cap = cv2.VideoCapture(camera_id) # TODO: cv2.CAP_DSHOW 能加速摄像头开启，但会导致视频保存出错？ # self.cap = cv2.VideoCapture( # camera_id) if camera_id == 0 else cv2.VideoCapture( # camera_id, cv2.CAP_DSHOW) # 调用外部摄像头需设置cv2.CAP_DSHOW self.is_camera = True else: self.cap = cv2.VideoCapture(input_video) self.is_camera = False # self.cap.setExceptionMode(True) # opencv 4.6 的自动旋转错误，采用自定义的旋转方式 # self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0.0) # self.rotate_code = self.check_rotation( # self.cap.get(cv2.CAP_PROP_ORIENTATION_META)) self.rotate_code = None self.t_start_save_video = None self.save_with_av = False self.out_stream = None self.container = None self.stream = None self.previous_pts = 0 def __del__(self): # self.cap.release() # logger.info('Camera(%s) closed.', self.__class__.__name__) # if self.out_stream: # self.out_stream.release() # if self.container and self.t_start_save_video: # self.release_container() self.close() def get_save_fps(self): return int(self.cap.get(cv2.CAP_PROP_FPS)) def open_camera(self): success = self.cap.open(self.camera_id) if success: logger.info('Open camera(%s) succeed.', self.__class__.__name__) else: logger.error('Open camera(%s) failed.', self.__class__.__name__) # if camera_id == 0: # self.cap.open(camera_id) # else: # self.cap.open(camera_id, cv2.CAP_DSHOW) def close(self, only_save: bool = False): """关闭摄像头与结束视频保存 如果only_save为true，则结束视频保存，但不关闭摄像头；否则关闭摄像头与结束视频保存 Args: only_save (bool, optional): 是否仅结束视频保存. Defaults to False. """ if not only_save: self.cap.release() logger.info('Camera(%s) closed.', self.__class__.__name__) if self.out_stream: self.out_stream.release() self.out_stream = None self.release_container() self.container = None def set_output_video(self, output_video, save_with_av=False): """ 设置输出视频 使用摄像头的情况下，必须在开摄像头之后调用，否则参数获取失败，无法正确设置输出视频 Args: output_video (string): 要保存的视频文件路径 save_with_av (bool, optional): 使用av库进行保存 """ self.save_with_av = save_with_av if not self.save_with_av: # video info # fourcc = int(self.cap.get(cv2.CAP_PROP_FOURCC)) # NOTICE: 这里需用 avc1 否则前端无法正常显示 fourcc = cv2.VideoWriter_fourcc(*'avc1') fps = self.get_save_fps() frame_size = (int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) # file to save video self.out_stream = cv2.VideoWriter(output_video, fourcc, fps, frame_size) else: assert self.is_camera,\ 'Do not save video with av when process recorded video!' self.container = av.open(output_video, mode='w') # NOTICE: 这里需使用 h264, 否则前端无法正常显示 self.stream = self.container.add_stream( 'h264', rate=int(self.cap.get(cv2.CAP_PROP_FPS))) # alibi frame rate self.stream.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.stream.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.stream.pix_fmt = 'yuv420p' self.stream.codec_context.time_base = Fraction( 1, int(self.cap.get(cv2.CAP_PROP_FPS))) def is_ok(self): if self.cap and self.cap.isOpened(): return True else: logger.debug('Camera not ready!!!') return False def check_rotation(self, rotate): rotate_code = None if int(rotate) == 270: rotate_code = cv2.ROTATE_90_CLOCKWISE elif int(rotate) == 180: rotate_code = cv2.ROTATE_180 elif int(rotate) == 90: rotate_code = cv2.ROTATE_90_COUNTERCLOCKWISE return rotate_code def correct_rotation(self, frame, rotate_code): return cv2.rotate(frame, rotate_code) def process(self, save=True): try: success, image = self.cap.read() if not success: logger.debug('Ignoring empty camera frame.') if self.rotate_code is not None: image = self.correct_rotation(image, self.rotate_code) except cv2.error as exc: logger.error( 'read data from camera(%s) failed, it may be disconnected: %s', self.__class__.__name__, exc) raise exc t_read = time.time() if success and save: self.save_video(image, t_read) return success, image def save_video(self, image, t_read): if self.save_with_av: self.save_video_with_av(image, t_read) else: self.save_video_with_opencv(image) def save_video_with_opencv(self, image): if not self.out_stream: return try: assert self.out_stream.isOpened(), 'Cannot open video for writing' self.out_stream.write(image) except Exception as exc: logger.error('Fail to save video %s: %s', self.out_stream, exc) def save_video_with_av(self, image, t_start): """Save video with [av](https://github.com/PyAV-Org/PyAV) Args: image (np.ndarray): frame to save t_start (float): timestamp of this frame """ if not self.container: return try: if not self.t_start_save_video: self.t_start_save_video = t_start frame = av.VideoFrame.from_ndarray(image, format='bgr24') # Presentation Time Stamp (seconds -> counts of time_base) delta_t = t_start - self.t_start_save_video if delta_t < 0.0: return pts = int(round(delta_t / self.stream.codec_context.time_base)) logger.debug('pts: %d', pts) if pts > self.previous_pts: frame.pts = pts self.previous_pts = frame.pts for packet in self.stream.encode(frame): self.container.mux(packet) except ValueError as exc: logger.debug('Fail to save frame of video %s: %s', self.container, exc) def release_container(self): if self.t_start_save_video: self.av_finish_with_a_blank_frame() # Close the file if self.container: self.container.close() self.t_start_save_video = None self.previous_pts = 0 def av_finish_with_a_blank_frame(self): # finish it with a blank frame, so the "last" frame actually gets # shown for some time this black frame will probably be shown for # 1/fps time at least, that is the analysis of ffprobe try: image = np.zeros((self.stream.height, self.stream.width, 3), dtype=np.uint8) frame = av.VideoFrame.from_ndarray(image, format='bgr24') pts = int( round((time.time() - self.t_start_save_video) / self.stream.codec_context.time_base)) logger.debug('last pts: %d', pts) frame.pts = pts if pts > self.previous_pts else self.previous_pts + 1 for packet in self.stream.encode(frame): self.container.mux(packet) # Flush stream for packet in self.stream.encode(): self.container.mux(packet) except ValueError as exc: logger.debug('Fail to save frame of video %s: %s', self.container, exc) def generator(self): while self.is_ok(): success, frame = self.process() # 使用generator函数输出视频流， 每次请求输出的content类型是image/jpeg if success: # 因为opencv读取的图片并非jpeg格式，因此要用motion JPEG模式需要先将图片转码成jpg格式图片 ret, jpeg = cv2.imencode('.jpg', frame) # t_end = time.time() # logger.debug("Time for process: %fs", t_end - t_start) yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n\r\n') def create_data_dir(subject_id, train_id): """为保存视频数据创建文件夹 Args: subject_id (_type_): _description_ train_id (_type_): _description_ """ path = f'{settings.DATA_PATH}/{subject_id}/{train_id}' try: os.makedirs(path) except OSError: logger.debug('Folder already exists!') return path def json_generator(feeder): while feeder.is_ok(): # time.sleep(1 / 30.0) success, _, data = feeder.process(only_keypoint=False) if success: json_data = json.dumps(data) yield f'data:{json_data}\n\n'