main.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. #!/usr/bin/env python3
  2. import asyncio
  3. import base64
  4. import concurrent.futures
  5. import signal
  6. import time
  7. from typing import Optional
  8. import cv2
  9. import numpy as np
  10. import psutil
  11. from fastapi import Response
  12. from icecream import ic
  13. from nicegui import app
  14. from nicegui import globals as nicegui_globals
  15. from nicegui import ui
  16. # we need two executors to schedule IO and CPU intensive tasks with loop.run_in_executor()
  17. process_pool_executor = concurrent.futures.ProcessPoolExecutor()
  18. thread_pool_executor = concurrent.futures.ThreadPoolExecutor()
  19. # in case you don't have a webcam, this will provide a black placeholder image
  20. black_1px = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAAXNSR0IArs4c6QAAAA1JREFUGFdjYGBg+A8AAQQBAHAgZQsAAAAASUVORK5CYII='
  21. placeholder = Response(content=base64.b64decode(black_1px.encode('ascii')), media_type='image/png')
  22. # OpenCV is used to access the webcam
  23. video_capture = cv2.VideoCapture(0)
  24. def convert(frame: np.ndarray) -> Optional[bytes]:
  25. if not frame_updater.active:
  26. return None
  27. _, imencode_image = cv2.imencode('.jpg', frame)
  28. return imencode_image.tobytes()
  29. @app.get('/video/frame')
  30. # thanks to FastAPI's "app.get" it is easy to create a web route which always provides the latest image from OpenCV
  31. async def grab_video_frame() -> Response:
  32. if not frame_updater.active:
  33. return placeholder
  34. loop = asyncio.get_running_loop()
  35. if not video_capture.isOpened():
  36. return placeholder
  37. # the video_capture.read call is a blocking function, so we run it in a separate thread it to avoid blocking the event loop
  38. _, frame = await loop.run_in_executor(thread_pool_executor, video_capture.read)
  39. if frame is None:
  40. return placeholder
  41. # "convert" is a cpu intensive function, so we run it in a separate process to avoid blocking the event loop and GIL
  42. jpeg = await loop.run_in_executor(process_pool_executor, convert, frame)
  43. if not jpeg:
  44. return placeholder
  45. return Response(content=jpeg, media_type='image/jpeg')
  46. # For non-flickering image updates an interactive image is much better than ui.image().
  47. video_image = ui.interactive_image().classes('w-full h-full')
  48. # A timer constantly updates the source of the image.
  49. # Because data from same paths are cached by the browser, we must force an update by adding the current timestamp to the source.
  50. frame_updater = ui.timer(interval=0.1, callback=lambda: video_image.set_source(f'/video/frame?{time.time()}'))
  51. def stop_updates(signum, frame):
  52. frame_updater.active = False
  53. ui.timer(1, lambda: signal.default_int_handler(signum, frame), once=True)
  54. async def cleanup():
  55. for client in nicegui_globals.clients.keys():
  56. await app.sio.disconnect(client)
  57. video_capture.release()
  58. thread_pool_executor.shutdown()
  59. # the process pool executor must be shutdown when the app is closed, otherwise the process will not exit
  60. process_pool_executor.shutdown()
  61. await asyncio.sleep(1)
  62. app.on_shutdown(cleanup)
  63. signal.signal(signal.SIGINT, stop_updates)
  64. ui.run()