main.py 3.5 KB

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