vcamera_server.py

import argparse
import os.path
from os import path as osp
from contextlib import nullcontext

from loguru import logger
try:
    import mvsdk
    from third_party.mvcam.vcamera import vCameraSystem
except ImportError:
    logger.warning('mvsdk not found, MindVision camera will not be supported')
import time
import tqdm
import cv2
from collections import deque
import threading
from fastapi import FastAPI
from fastapi.responses import Response
import uvicorn
from ImplicitRDP.common.precise_sleep import precise_sleep
from ImplicitRDP.real_world.simple_camera.simple_realsense_camera import SimpleRealsenseCamera
from ImplicitRDP.real_world.simple_camera.simple_usb_camera import SimpleUSBCamera


class vCameraSever:
    def __init__(self, host_ip, port, camera_type="RealSense", camera_id=0, fps=30, buffer_size=60, video_size=(1280, 800)):
        self.camera_type = camera_type
        self.camera_id = camera_id

        self.host_ip = host_ip
        self.port = port

        self.fps = fps
        self.capture_interval_time = 1.0 / fps
        self.capture_buffer = deque(maxlen=buffer_size)

        self.camera = self.get_camera(camera_id)

        self.app = FastAPI()
        self.setup_routes()

        self.video_size = video_size
        self.video_writer = None
        self.record_event = threading.Event()
        self.stop_event = threading.Event()
        self.mutex = threading.Lock()
    
    def get_camera(self, camera_id):
        if self.camera_type == "RealSense":
            camera_id = str(camera_id)
            camera = SimpleRealsenseCamera(camera_serial_number=camera_id,
                                           rgb_resolution=(1920, 1080),
                                           exposure=480,
                                           fps=self.fps)
        elif self.camera_type == "MindVision":
            camera_id = int(camera_id)
            if 'mvsdk' not in globals():
                raise ValueError('mvsdk not found, MindVision camera will not be supported')
            DevList = mvsdk.CameraEnumerateDevice()
            if len(DevList) - 1 < camera_id:
                camera = None
            else:
                cam_sys = vCameraSystem()
                camera = cam_sys[camera_id]

            if camera is None:
                raise ValueError('No Mindvision camera found!')
        elif self.camera_type == "USB":
            camera_id = int(camera_id)
            camera = SimpleUSBCamera(camera_index=camera_id,
                                    width=640,
                                    height=360,
                                    fps=self.fps,
                                    exposure=-6,
                                    contrast=100)
        else:
            raise ValueError(f'Unknown camera type: {self.camera_type}')
    
        return camera

    def setup_routes(self):
        @self.app.get('/get_capture')
        async def get_capture():
            with self.mutex:
                capture_buffer_size = len(self.capture_buffer)
                if capture_buffer_size == 0:
                    capture = None
                else:
                    capture = self.capture_buffer.popleft()
            logger.info(f'Rest capture buffer size: {capture_buffer_size}')
            return Response(content=capture, media_type="application/octet-stream")

        @self.app.get('/peek_latest_capture')
        async def peek_latest_capture():
            with self.mutex:
                capture_buffer_size = len(self.capture_buffer)
                if capture_buffer_size == 0:
                    capture = None
                else:
                    capture = self.capture_buffer[-1]
            return Response(content=capture, media_type="application/octet-stream")

        @self.app.post('/start_recording/{video_path:path}')
        async def start_recording(video_path):
            if not osp.exists(osp.dirname(video_path)):
                os.makedirs(osp.dirname(video_path))
            video_path = ".".join(video_path.split('.')[:-1]) + f'_camera{self.camera_id}.mp4'
            if osp.exists(video_path):
                video_path = ".".join(video_path.split('.')[:-1]) + f'{time.strftime("_%Y%m%d_%H%M%S")}.mp4'
                logger.warning(f'Video path already exists, save to {video_path}')
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            with self.mutex:
                self.video_writer = cv2.VideoWriter(video_path, fourcc, self.fps, self.video_size)
            self.record_event.set()
            logger.info(f'Start recording to {video_path}')
            return Response(content='Start recording', media_type="text/plain")

        @self.app.post('/stop_recording')
        async def stop_recording():
            self.record_event.clear()
            with self.mutex:
                if self.video_writer is not None:
                    self.video_writer.release()
                    self.video_writer = None
            logger.info(f'Stop recording')
            return Response(content='Stop recording', media_type="text/plain")

    def run(self):
        capture_thread = threading.Thread(target=self.capture, args=(self.stop_event,), daemon=True)
        try:
            capture_thread.start()
            uvicorn.run(self.app, host=self.host_ip, port=self.port)
        except Exception as e:
            logger.exception(e)
            raise e
        finally:
            self.stop_event.set()
            capture_thread.join()

    def capture(self, stop_event):
        if self.camera_type == "RealSense":
            context = nullcontext()
        elif self.camera_type == "MindVision":
            context = self.camera
        elif self.camera_type == "USB":
            context = nullcontext()
            self.camera.start()
        else:
            raise ValueError(f'Unknown camera type: {self.camera_type}')
            
        try:
            with context as c:
                logger.info(f'Start camera capture')
                with tqdm.tqdm(range(1)) as pbar:
                    start_t = time.time()
                    cnt = 0
                    while not stop_event.is_set():
                        capture_start_t = time.time()
                        cnt += 1
                        if self.camera_type == "RealSense":
                            img, _ = self.camera.get_frames()
                        elif self.camera_type == "MindVision":
                            img = c.read()
                            img = img[:, :, ::-1]  # RGB -> BGR
                        elif self.camera_type == "USB":
                            img = self.camera.get_rgb_frame()
                        else:
                            raise ValueError(f'Unknown camera type: {self.camera_type}')

                        video_img = cv2.resize(img, self.video_size)
                        with self.mutex:
                            if self.record_event.is_set():
                                self.video_writer.write(video_img)

                        img = cv2.imencode('.jpg', img)[1].tobytes()
                        with self.mutex:
                            self.capture_buffer.append(img)
                        cur_time = time.time()
                        precise_sleep(max(0., self.capture_interval_time - (cur_time - capture_start_t)))
                        pbar.set_description(f"fps={cnt / (cur_time - start_t)}")
        finally:
            if self.camera_type == "USB":
                self.camera.stop()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--host_ip', type=str, default='192.168.2.187')
    parser.add_argument('--port', type=int, default=8102)
    parser.add_argument('--camera_type', type=str, default="RealSense", choices=["RealSense", "MindVision", "USB"])
    parser.add_argument('--camera_id', type=str, default="036422060422")
    parser.add_argument('--fps', type=int, default=30)
    parser.add_argument('--video_size', type=str, default='1920,1080')
    args = parser.parse_args()

    vcamera_server = vCameraSever(host_ip=args.host_ip,
                                  port=args.port,
                                  camera_type=args.camera_type,
                                  camera_id=args.camera_id,
                                  fps=args.fps,
                                  video_size=tuple(map(int, args.video_size.split(','))))
    vcamera_server.run()