main.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. #!/usr/bin/env python3
  2. import asyncio
  3. import base64
  4. import concurrent.futures
  5. import signal
  6. import time
  7. import cv2
  8. import numpy as np
  9. from fastapi import Response
  10. import nicegui.globals
  11. from nicegui import app, ui
  12. # We need an executor to schedule CPU-intensive tasks with `loop.run_in_executor()`.
  13. process_pool_executor = concurrent.futures.ProcessPoolExecutor()
  14. # In case you don't have a webcam, this will provide a black placeholder image.
  15. black_1px = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAAXNSR0IArs4c6QAAAA1JREFUGFdjYGBg+A8AAQQBAHAgZQsAAAAASUVORK5CYII='
  16. placeholder = Response(content=base64.b64decode(black_1px.encode('ascii')), media_type='image/png')
  17. # OpenCV is used to access the webcam.
  18. video_capture = cv2.VideoCapture(0)
  19. def convert(frame: np.ndarray) -> bytes:
  20. _, imencode_image = cv2.imencode('.jpg', frame)
  21. return imencode_image.tobytes()
  22. @app.get('/video/frame')
  23. # Thanks to FastAPI's `app.get`` it is easy to create a web route which always provides the latest image from OpenCV.
  24. async def grab_video_frame() -> Response:
  25. if not video_capture.isOpened():
  26. return placeholder
  27. loop = asyncio.get_running_loop()
  28. # The `video_capture.read` call is a blocking function.
  29. # So we run it in a separate thread (default executor) to avoid blocking the event loop.
  30. _, frame = await loop.run_in_executor(None, video_capture.read)
  31. if frame is None:
  32. return placeholder
  33. # `convert` is a CPU-intensive function, so we run it in a separate process to avoid blocking the event loop and GIL.
  34. jpeg = await loop.run_in_executor(process_pool_executor, convert, frame)
  35. return Response(content=jpeg, media_type='image/jpeg')
  36. # For non-flickering image updates an interactive image is much better than `ui.image()`.
  37. video_image = ui.interactive_image().classes('w-full h-full')
  38. # A timer constantly updates the source of the image.
  39. # Because data from same paths are cached by the browser,
  40. # we must force an update by adding the current timestamp to the source.
  41. ui.timer(interval=0.1, callback=lambda: video_image.set_source(f'/video/frame?{time.time()}'))
  42. async def disconnect() -> None:
  43. """Disconnect all clients from current running server."""
  44. for client in nicegui.globals.clients.keys():
  45. await app.sio.disconnect(client)
  46. def handle_sigint(signum, frame) -> None:
  47. # `disconnect` is async, so it must be called from the event loop; we use `ui.timer` to do so.
  48. ui.timer(0.1, disconnect, once=True)
  49. # Delay the default handler to allow the disconnect to complete.
  50. ui.timer(1, lambda: signal.default_int_handler(signum, frame), once=True)
  51. async def cleanup() -> None:
  52. # This prevents ugly stack traces when auto-reloading on code change,
  53. # because otherwise disconnected clients try to reconnect to the newly started server.
  54. await disconnect()
  55. # Release the webcam hardware so it can be used by other applications again.
  56. video_capture.release()
  57. # The process pool executor must be shutdown when the app is closed, otherwise the process will not exit.
  58. process_pool_executor.shutdown()
  59. app.on_shutdown(cleanup)
  60. # We also need to disconnect clients when the app is stopped with Ctrl+C,
  61. # because otherwise they will keep requesting images which lead to unfinished subprocesses blocking the shutdown.
  62. signal.signal(signal.SIGINT, handle_sigint)
  63. ui.run()