diff --git a/aruco_realtime.py b/aruco_realtime.py new file mode 100644 index 0000000..7eaaab7 --- /dev/null +++ b/aruco_realtime.py @@ -0,0 +1,46 @@ +import sys + +import cv2 +from cv2 import aruco + +# Define the ArUco dictionary +aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250) +aruco_params = cv2.aruco.DetectorParameters() + +# Specify the RTSP URL +rtsp_url = sys.argv[1] + +# Capture the video stream from the RTSP URL +cap = cv2.VideoCapture(rtsp_url) + +while True: + ret, frame = cap.read() + + # Check if frame is not empty + if frame is None: + print("Empty frame, exiting...") + break + + # Detect ArUco markers + gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) + corners, ids, rejectedImgPoints = aruco.detectMarkers( + gray, aruco_dict, parameters=aruco_params + ) + + # If some markers are found, draw them on the frame + if ids is not None: + frame = cv2.aruco.drawDetectedMarkers(frame, corners, ids) + + # Scale frame in half + frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5) + + # Display the frame + cv2.imshow("Frame", frame) + + # Wait for the 'q' key to exit the loop + if cv2.waitKey(1) & 0xFF == ord("q"): + break + +# Release the capture object and close all OpenCV windows +cap.release() +cv2.destroyAllWindows() diff --git a/bot.py b/bot.py index a7cbd90..7edaba1 100644 --- a/bot.py +++ b/bot.py @@ -1,13 +1,18 @@ import asyncio import os +import time import traceback import warnings -from typing import cast, Optional +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime +from typing import cast, List import aiohttp import cv2 +import numpy as np import telegram from aiohttp import BasicAuth +from cv2 import aruco from pytapo import Tapo from telegram import Update, Message from telegram.ext import Updater @@ -32,6 +37,7 @@ class Bot: self.me = None self.tapo = Tapo(self.camera_ip, self.camera_user, self.camera_password) + self.executor = ThreadPoolExecutor(max_workers=len(os.sched_getaffinity(0))) def _get_presets(self): presets = self.tapo.getPresets() @@ -56,14 +62,133 @@ class Bot: ) as resp: return await resp.text() - def _take_photo(self) -> Optional[bytes]: + def _take_photo_blocking( + self, + adjust_perspective=True, + timeout: float = 5.0, + ) -> List[cv2.typing.MatLike]: + # Prepare the camera + print("Disabling privacy mode and setting auto day/night mode...") + self.tapo.setPrivacyMode(False) + self.tapo.setDayNightMode("auto") + time.sleep(1) + + # Take the color image vcap = cv2.VideoCapture( f"rtsp://{self.camera_user}:{self.camera_password}@{self.camera_ip}:554/stream1" ) - ret, frame = vcap.read() - if not ret: - return None - return frame + print("Taking color image...") + ret, pretty_image = vcap.read() + while pretty_image is None: + ret, pretty_image = vcap.read() + + self.tapo.setDayNightMode("on") + + if not adjust_perspective: + return [pretty_image] + + # Iterate until we find all 4 aruco markers or timeout + aruco_corners = {} + annotated_image = None + aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_6X6_250) + aruco_params = aruco.DetectorParameters() + + t0 = time.time() + print("Taking image with ArUco markers...") + while len(aruco_corners) < 4: + if time.time() - t0 > timeout: + print( + "Timeout waiting for ArUco markers, returning only original image" + ) + return [pretty_image] + + ret, annotated_image = vcap.read() + if not ret: + continue + + # Detect the markers + corners, ids, rejected = aruco.detectMarkers( + annotated_image, aruco_dict, parameters=aruco_params + ) + for corner, i in zip(corners, ids): + aruco_corners[i[0]] = corner + assert annotated_image is not None + + del vcap + + print( + "Found all ArUco markers, enabled privacy mode and set auto day/night mode..." + ) + self.tapo.setDayNightMode("auto") + self.tapo.setPrivacyMode(True) + + corners = [aruco_corners[i] for i in range(1, 5)] + ids = np.array([[i] for i in range(1, 5)]) + aruco.drawDetectedMarkers(annotated_image, corners, ids) + + # Annotate the image with the detected markers and apply the perspective transform to the pretty image + + # Get the outermost points of each marker + bl_marker = corners[0].squeeze()[3] # bottom left marker + tl_marker = corners[1].squeeze()[0] # top left marker + tr_marker = corners[2].squeeze()[1] # top right marker + bc_marker = corners[3].squeeze()[3] # bottom center marker + + # Calculate the fourth point by computing the line through the bottom markers and intersecting with the vertical + # line through the top right marker + slope = (bc_marker[1] - bl_marker[1]) / (bc_marker[0] - bl_marker[0]) + y_intersection = slope * (tr_marker[0] - bc_marker[0]) + bc_marker[1] + fourth_point = [tr_marker[0], y_intersection] + + rectangle_points = np.array( + [tl_marker, bl_marker, fourth_point, tr_marker], dtype="float32" + ) + + # Expand the rectangle slightly + centroid = np.mean(rectangle_points, axis=0) + expanded_rectangle_points = ( + rectangle_points + (rectangle_points - centroid) * 0.025 + ) + + # Draw the expanded rectangle on the annotated image + cv2.polylines( + annotated_image, + [np.int32(expanded_rectangle_points)], + isClosed=True, + color=(0, 255, 0), + thickness=3, + ) + + # Define destination points for perspective transform, maintaining a 3:2 aspect ratio + width, height = 300 * 5, 200 * 5 + dst_pts = np.array( + [[0, 0], [0, height], [width, height], [width, 0]], dtype="float32" + ) + + matrix = cv2.getPerspectiveTransform(expanded_rectangle_points, dst_pts) + warped = cv2.warpPerspective(pretty_image, matrix, (width, height)) + + return [warped, annotated_image] + + async def _photo_command(self, chat_id: int, adjust_perspective: bool = True): + await self.bot.send_chat_action(chat_id=chat_id, action="upload_photo") + + photos = await self.take_photo(adjust_perspective=adjust_perspective) + jpegs = [cv2.imencode(".jpg", photo)[1].tobytes() for photo in photos] + + media = [ + telegram.InputMediaPhoto( + media=telegram.InputFile(jpeg, filename=f"photo{i}.jpg", attach=True), + filename=f"photo{i}.jpg", + ) + for i, jpeg in enumerate(jpegs) + ] + + await self.bot.send_media_group( + chat_id=chat_id, + media=media, + caption=str(datetime.now().strftime("%A, %B %d, %Y %H:%M:%S")), + ) async def parse_message(self, msg: Message): match msg.text: @@ -113,40 +238,26 @@ class Bot: text=f"Light is {state}", ) case "/photo": - await self.bot.send_chat_action( - chat_id=msg.chat_id, action="upload_photo" - ) + await self._photo_command(msg.chat_id, adjust_perspective=True) + case "/photo_unprocessed": + await self._photo_command(msg.chat_id, adjust_perspective=False) - light_state = await self._get_item_state() - if light_state == "OFF": - print("Turning light on") - await self._send_item_command("ON") + async def take_photo( + self, adjust_perspective=True, timeout=5.0 + ) -> List[cv2.typing.MatLike]: + item_state = await self._get_item_state() + if item_state == "OFF": + print("Turning light on") + await self._send_item_command("ON") - print("Disabling privacy mode") - self.tapo.setPrivacyMode(False) - await asyncio.sleep(2) - print("Taking photo") - photo = self._take_photo() - print("Enabling privacy mode") - self.tapo.setPrivacyMode(True) - - if light_state == "OFF": - print("Turning light back off") - await self._send_item_command("OFF") - - if photo is None: - await self.bot.send_message( - chat_id=msg.chat_id, - text=f"Error taking photo", - ) - return - # Encode photo to jpeg - photo = cv2.imencode(".jpg", photo)[1].tobytes() - - await self.bot.send_photo( - chat_id=msg.chat_id, - photo=telegram.InputFile(photo, filename="photo.jpg"), - ) + try: + return await asyncio.get_event_loop().run_in_executor( + self.executor, self._take_photo_blocking, adjust_perspective, timeout + ) + finally: + if item_state == "OFF": + print("Turning light back off") + await self._send_item_command("OFF") async def run(self): async with self.bot: @@ -155,6 +266,7 @@ class Bot: updater = Updater(bot=self.bot, update_queue=asyncio.Queue()) async with updater: queue = await updater.start_polling(allowed_updates=[Update.MESSAGE]) + print("Bot is up and running") while True: # noinspection PyBroadException @@ -164,7 +276,16 @@ class Bot: if not upd.message or upd.message.chat_id != self.chat_id: print("Ignoring message") continue - await self.parse_message(upd.message) + # noinspection PyBroadException + try: + await self.parse_message(upd.message) + except Exception: + traceback.print_exc() + exc = traceback.format_exc() + await self.bot.send_message( + chat_id=self.chat_id, + text=f"Error: {exc}", + ) except Exception: traceback.print_exc() diff --git a/commands.txt b/commands.txt index 3265317..caef359 100644 --- a/commands.txt +++ b/commands.txt @@ -4,3 +4,4 @@ reposition - Re-aim the camera light_on - Turn on the light light_off - Turn off the light light_status - Check the status of the light +photo_unprocessed - Take a photo of the grocery list board without processing \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index a3852d6..29453a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ pytapo python-telegram-bot opencv-python-headless +opencv-contrib-python-headless aiohttp \ No newline at end of file