{"cells":[{"cell_type":"code","execution_count":1,"metadata":{},"outputs":[],"source":["# Import Libraries\n","import asyncio\n","import cv2\n","import numpy as np\n","import math\n","from ultralytics import YOLO\n","\n","# Load Models\n","obb_model = YOLO(\"model/obbModel.pt\")\n","pose_model = YOLO(\"model/poseModel.pt\")\n","\n","model = [obb_model, pose_model]\n","\n","# Line Angle Measurement\n","async def line_angle(x1, y1, x2, y2):\n"," dx = x2 - x1\n"," dy = y2 - y1\n"," return math.atan2(dy, dx)\n","\n","\n","\n","async def measureObject(img, model):\n"," obb_predict = await asyncio.get_event_loop().run_in_executor(None, model[0].predict, img)\n"," pose_predict = await asyncio.get_event_loop().run_in_executor(None, model[1].predict, img)\n","\n"," obb_points = obb_predict[0].obb.xyxyxyxy[0].cpu().numpy().astype(int)\n"," pose_points = pose_predict[0].keypoints.xy[0].cpu().numpy().astype(int)\n","\n"," xmidtop = int((obb_points[0][0] + obb_points[1][0]) / 2)\n"," ymidtop = int((obb_points[0][1] + obb_points[1][1]) / 2)\n"," xmidbot = int((obb_points[2][0] + obb_points[3][0]) / 2)\n"," ymidbot = int((obb_points[2][1] + obb_points[3][1]) / 2)\n","\n"," xmidmid = int((xmidtop + xmidbot) / 2)\n"," ymidmid = int((ymidtop + ymidbot) / 2)\n"," img_with_lines = img.copy()\n","\n"," cv2.line(img_with_lines, (xmidtop, ymidtop), (xmidbot, ymidbot), (0, 255, 0), 2) \n"," cv2.line(img_with_lines, (xmidmid, ymidmid), (xmidtop, ymidtop), (255, 0, 0), 2) \n"," cv2.line(img_with_lines, (xmidmid, ymidmid), (xmidbot, ymidbot), (255, 0, 0), 2) \n","\n"," p1 = int(abs(pose_points[0][0] - obb_points[0][0]))\n"," q1 = int(abs(pose_points[0][1] - obb_points[0][1]))\n"," p2 = int(abs(pose_points[0][0] - obb_points[3][0]))\n"," q2 = int(abs(pose_points[0][1] - obb_points[3][1]))\n"," d1 = p1 + q1\n"," d2 = p2 + q2\n","\n"," if d1 < d2:\n"," cx = xmidtop\n"," cy = ymidtop\n"," else:\n"," cx = xmidbot\n"," cy = ymidbot\n","\n"," angle_rad = await line_angle(xmidmid, ymidmid, cx, cy)\n","\n"," angle = -1 * (math.degrees(angle_rad))\n","\n"," if angle < 0:\n"," angle += 360\n"," \n"," cv2.putText(img_with_lines, f'Angle: {angle:.2f}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)\n","\n"," return img_with_lines, {\"Center\": {\"X\": xmidmid, \"Y\": ymidmid}, \"Angle\": \"{:.2f}\".format(angle)}"]},{"cell_type":"code","execution_count":4,"metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["{'Center': {'X': 298, 'Y': 480}, 'Angle': '135.58'}\n"]}],"source":["# Testing\n","img = cv2.imread(\"img/3.png\") # Change the image and try\n","\n","small_img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5)\n","img_with_lines, center_and_angle_info = await measureObject(small_img, model)\n","center_x, center_y = center_and_angle_info[\"Center\"][\"X\"], center_and_angle_info[\"Center\"][\"Y\"]\n","\n","cv2.circle(img_with_lines, (center_x, center_y), 5, (255, 0, 255), -1)\n","cv2.putText(img_with_lines, f'x: {center_x}', (center_x + 10, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)\n","cv2.putText(img_with_lines, f'y: {center_y}', (center_x, center_y + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)\n","\n","cv2.imshow(\"Image with lines\", img_with_lines)\n","cv2.waitKey(0)\n","cv2.destroyAllWindows()\n","\n","print(center_and_angle_info)"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":[]}],"metadata":{"colab":{"authorship_tag":"ABX9TyO8vuV19KedxeHeSC6oPsS9","provenance":[]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.3"}},"nbformat":4,"nbformat_minor":0}