diff --git a/autopilot.py b/autopilot.py index 07d7313..3422816 100644 --- a/autopilot.py +++ b/autopilot.py @@ -55,6 +55,7 @@ class AutoPilot(Pilot): self.frame_count = 0 self.image_center = (0, 0) # Будет обновлено при получении первого изображения self.vis_manager = viz_manager # Менеджер визуализации + self.reserved_pos = None # Пороговые значения качества сопоставления/гомографии self.min_inliers: int = 12 @@ -65,7 +66,7 @@ class AutoPilot(Pilot): # Инициализация ORB детектора self.orb_detector = cv2.ORB_create( - nfeatures=300, + nfeatures=1800, scaleFactor=1.3, nlevels=4, edgeThreshold=19, @@ -159,7 +160,7 @@ class AutoPilot(Pilot): """ # Используем RANSAC для оценки матрицы гомографии H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.0) - RH, rmask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 1.0) + # RH, rmask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 1.0) if H is None: return None @@ -214,26 +215,34 @@ class AutoPilot(Pilot): 'inliers_ratio': inlier_ratio, 'rmse': rmse } + + + + def calc_position( + self, + mat: np.ndarray, + x_source: float, + y_source: float, + angle_source: float + ) -> tuple[float, float, float] | None: + """ Возвращает новые координаты и поворот на основе матрицы преобразования """ - def update_drone_position(self, transformation_info: dict): - """ - Обновляет позицию и угол БПЛА на основе трансформации изображения - Координаты уже отцентрированы относительно центра изображения - """ - tx, ty = transformation_info['translation'] - rotation = transformation_info['rotation'] - scale = transformation_info['scale'] + tx, ty = -mat[0, 2], -mat[1, 2] + + # Вычисляем угол поворота + rotation = -np.arctan2(mat[1, 0], mat[0, 0]) + print("[HOMOGRAPHY]", mat) + print("[ROTATION]", rotation) # Координаты уже отцентрированы, поэтому используем их напрямую - dx_meters = ty - dy_meters = tx + dx_meters = tx + dy_meters = ty - # Обновляем угол БПЛА - self.angle += rotation + angle_global = angle_source + rotation # Применяем поворот к смещению (учитываем текущий угол БПЛА) - cos_angle = math.cos(self.angle) - sin_angle = math.sin(self.angle) + cos_angle = math.cos(angle_global) + sin_angle = math.sin(angle_global) # Поворачиваем смещение в глобальные координаты # Обратите внимание: dy_meters инвертирован, так как в изображениях Y направлен вниз @@ -241,11 +250,30 @@ class AutoPilot(Pilot): dy_global = dx_meters * sin_angle + dy_meters * cos_angle # Обновляем координаты БПЛА - self.x += dx_global - self.y += dy_global + x = x_source + dx_global + y = y_source + dy_global # Нормализуем угол в диапазоне [-π, π] - self.angle = math.atan2(math.sin(self.angle), math.cos(self.angle)) + angle = math.atan2(math.sin(angle_global), math.cos(angle_global)) + + return x, y, angle + + def update_drone_position(self, transformation_info: dict): + """ + Обновляет позицию и угол БПЛА на основе трансформации изображения + Координаты уже отцентрированы относительно центра изображения + """ + homography = transformation_info['homography'] + + self.x, self.y, self.angle = self.calc_position( + homography, + self.x, self.y, self.angle + ) + + if self.reserved_pos is not None: + self.reserved_pos = self.calc_position( + homography, *self.reserved_pos + ) def get_drone_state(self) -> dict: """ @@ -258,15 +286,17 @@ class AutoPilot(Pilot): 'angle_degrees': math.degrees(self.angle), 'frame_count': self.frame_count } + + def get_position_by_chunk(self) -> tuple[float, float, float] | None: # Пытаемся найти ориентир на картинке: current_image = self.prev_image landmark_image = cv2.imread(Path('chunks') / f'chunk_{self.target_idx}.png', cv2.IMREAD_COLOR_RGB) - src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(current_image, landmark_image) + src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(landmark_image, current_image) if src_pts is not None and dst_pts is not None and self.vis_manager: - self.vis_manager.update_chunk_matches(current_image, landmark_image, kp1, kp2, matches) + self.vis_manager.update_chunk_matches(landmark_image, current_image, kp1, kp2, matches) if src_pts is not None and dst_pts is not None: # Оцениваем матрицу трансформации @@ -280,26 +310,19 @@ class AutoPilot(Pilot): rmse = landmark_transform.get('rmse', None) ok_rmse = (rmse is not None and rmse <= self.max_reproj_rmse) if ok_scale and ok_inliers and ok_ratio and ok_rmse: + print("[HELP]") + print("Matrix", landmark_transform['homography']) + print("Position", self.x, self.y) + print("Position of point", self.points[self.target_idx]) print("[PILOT]", rmse, ratio, ok_rmse) # if False: # Считаем абсолютную позу относительно координат ориентира landmark_world_x, landmark_world_y = self.points[self.target_idx] - tx, ty = landmark_transform['translation'] - # Смещение в системе БПЛА: (dx_meters, dy_meters) - dx_meters = ty - dy_meters = tx - # Используем оценённый абсолютный угол из гомографии относительно карты (север-вверх) - absolute_angle = landmark_transform['rotation'] - cos_a = math.cos(absolute_angle) - sin_a = math.sin(absolute_angle) - # Переход в глобальные координаты карты - dx_global = dx_meters * cos_a - dy_meters * sin_a - dy_global = dx_meters * sin_a + dy_meters * cos_a - x = landmark_world_x + dx_global - y = landmark_world_y + dy_global - angle = math.atan2(math.sin(absolute_angle), math.cos(absolute_angle)) + landmark_angle = 0 + homography = landmark_transform['homography'] + # homography = np.linalg.inv(homography) print(f" [Pilot] Landmark correction applied (inliers={landmark_transform['inliers']}, ratio={ratio:.2f}, rmse={rmse:.2f})") - return (x, y, angle) + return self.calc_position(homography, landmark_world_x, landmark_world_y, landmark_angle) return None @@ -372,45 +395,12 @@ class AutoPilot(Pilot): grid_step=85) # Пытаемся найти ориентир на картинке: - landmark_image = cv2.imread(Path('chunks') / f'chunk_{self.target_idx}.png', cv2.IMREAD_COLOR_RGB) - src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(current_image, landmark_image) + self.prev_image = current_image + npos = self.get_position_by_chunk() + if npos is not None: + self.reserved_pos = npos - if src_pts is not None and dst_pts is not None and self.vis_manager: - self.vis_manager.update_chunk_matches(current_image, landmark_image, kp1, kp2, matches) - - # if src_pts is not None and dst_pts is not None: - # # Оцениваем матрицу трансформации - # landmark_transform = self.estimate_transformation_matrix(src_pts, dst_pts) - # # Если ориентир достоверно найден — скорректируем глобальные координаты и угол - # if landmark_transform is not None: - # ok_scale = (self.min_scale <= landmark_transform['scale'] <= self.max_scale) - # ok_inliers = (landmark_transform.get('inliers', 0) >= self.min_inliers) - # ratio = landmark_transform.get('inliers_ratio', 0.0) - # ok_ratio = (ratio >= self.min_inlier_ratio) - # rmse = landmark_transform.get('rmse', None) - # ok_rmse = (rmse is not None and rmse <= self.max_reproj_rmse) - # if ok_scale and ok_inliers and ok_ratio and ok_rmse: - # print("[PILOT]", rmse, ratio, ok_rmse) - # # if False: - # # Считаем абсолютную позу относительно координат ориентира - # landmark_world_x, landmark_world_y = self.points[self.target_idx] - # tx, ty = landmark_transform['translation'] - # # Смещение в системе БПЛА: (dx_meters, dy_meters) - # dx_meters = ty - # dy_meters = tx - # # Используем оценённый абсолютный угол из гомографии относительно карты (север-вверх) - # absolute_angle = landmark_transform['rotation'] - # cos_a = math.cos(absolute_angle) - # sin_a = math.sin(absolute_angle) - # # Переход в глобальные координаты карты - # dx_global = dx_meters * cos_a - dy_meters * sin_a - # dy_global = dx_meters * sin_a + dy_meters * cos_a - # self.x = landmark_world_x + dx_global - # self.y = landmark_world_y + dy_global - # self.angle = math.atan2(math.sin(absolute_angle), math.cos(absolute_angle)) - # print(f" [Pilot] Landmark correction applied (inliers={landmark_transform['inliers']}, ratio={ratio:.2f}, rmse={rmse:.2f})") - - if distance_to_target < 75: + if distance_to_target < 35: self.target_idx += 1 if self.target_idx == len(self.points): @@ -421,19 +411,28 @@ class AutoPilot(Pilot): (self.points[self.target_idx][1] - self.y) ** 2 ) + if self.reserved_pos is not None: + self.x, self.y, self.angle = self.reserved_pos + self.reserved_pos = None + # Вычисляем угол к цели target_angle = math.atan2(self.points[self.target_idx][1] - self.y, self.points[self.target_idx][0] - self.x) + angle_trajectory = self.angle + math.pi / 2 + + print("[ANGLE]", angle_trajectory, "->", target_angle) + # Вычисляем разность углов (направление поворота) - angle_diff = target_angle - self.angle + angle_diff = target_angle - angle_trajectory # Нормализуем разность углов в диапазон [-π, π] angle_diff %= 2 * math.pi if angle_diff >= math.pi: angle_diff -= 2 * math.pi - self.prev_image = current_image - return PilotCommand(max(min(0.6, angle_diff), -0.6), min(50., distance_to_target / 2)) + d_r = max(10, min(75., distance_to_target / 2)) + d_a_limit = d_r / 10 * 0.07 + return PilotCommand(max(min(d_a_limit, angle_diff), -d_a_limit), d_r) def reset_position(self, x: float = 0.0, y: float = 0.0, angle: float = 0.0): """Сбрасывает позицию и угол БПЛА""" diff --git a/main.py b/main.py index ae31dfa..763fb87 100644 --- a/main.py +++ b/main.py @@ -29,10 +29,10 @@ def main(): # Получаем траекторию от пользователя points = get_trajectory_points('map.jpg') - print(points) # points = [np.float64(0.5443937502226799), np.float64(0.4030424838774785)], [np.float64(0.18517133490120316), np.float64(0.4586935604608052)], [np.float64(0.1641887171838272), np.float64(0.5586383510594329)], [np.float64(0.587198290366127), # np.float64(0.5699957136274587)] - # points = [[np.float64(0.4931961629922826), np.float64(0.4439289891223716)], [np.float64(0.5099822571661834), np.float64(0.44165751660876645)], [np.float64(0.5083036477487933), np.float64(0.4564220879472001)], [np.float64(0.525089741922694), np.float64(0.4586935604608052)]] + # points = [[np.float64(0.5028362002950056), np.float64(0.508463417020251)], [np.float64(0.5079239482442035), np.float64(0.508463417020251)], [np.float64(0.5045321162780716), np.float64(0.5210964284169014)], [np.float64(0.509054558899581), np.float64(0.5200436774671806)]] + # print(points) # Для каждой точки сделаем приближенный снимок yandexMap = YandexMap() @@ -105,6 +105,7 @@ def main(): vis_manager.set_target_points(points_coords) for i in range(10000000000): + print(f"Image #{i}") photo = simulator.get_photo() command = pilot.handle(photo) diff --git a/simulator.py b/simulator.py index c4251db..94d1b73 100644 --- a/simulator.py +++ b/simulator.py @@ -88,8 +88,8 @@ class Simulator: self.angle += dangle print(f" [Simulator] angle: {self.angle / math.pi * 180:.1f}°") velocity = max(velocity, 10) - dx = math.cos(self.angle) * velocity - dy = math.sin(self.angle) * velocity + dx = math.cos(math.pi / 2 + self.angle) * velocity + dy = math.sin(math.pi / 2 + self.angle) * velocity print(" [Simulator] dx, dy:", [dx, dy]) self.update_trajectory(dx, dy) action.move_by_offset(-dx, dy) @@ -102,7 +102,7 @@ class Simulator: im = Image.open(BytesIO(png)) # Применяем поворот как будто съемка с дрона - rotated_im = self.rotate_image_like_drone(im, math.pi / 2 - self.angle) + rotated_im = self.rotate_image_like_drone(im, -self.angle) return rotated_im def loop(self): @@ -138,8 +138,8 @@ class Simulator: action.click_and_hold() self.angle += dangle - dx = math.cos(self.angle) * velocity - dy = math.sin(self.angle) * velocity + dx = math.cos(math.pi / 2 + self.angle) * velocity + dy = math.sin(math.pi / 2 + self.angle) * velocity action.move_by_offset(-dx, dy) action.release() action.perform() diff --git a/test_autopilot.ipynb b/test_autopilot.ipynb index 8148a09..e968159 100644 --- a/test_autopilot.ipynb +++ b/test_autopilot.ipynb @@ -40,7 +40,7 @@ "\n", "vzm.set_target_points(points_coords)\n", "autopilot = AutoPilot(points_coords, [], vzm)\n", - "imgs = [Image.open(Path('images') / f'photo_{i}.png') for i in range(10)]\n", + "imgs = [Image.open(Path('images') / f'photo_{i}.png') for i in range(30)]\n", "autopilot.handle(imgs[0])\n", "vzm.update_drone_trajectory(autopilot.x, autopilot.y)\n" ] @@ -49,19 +49,54 @@ "cell_type": "code", "execution_count": 3, "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "-84.81833939223279\n", + "[[[ 0.43529349 -148.93070048]]]\n", + "164.0668026014628 -13.040291124133018 90.0\n" + ] + } + ], + "source": [ + "import cv2\n", + "import math\n", + "mat = np.array([[ 9.03878625e-02, 9.95998304e-01, -1.36910620e+01],\n", + " [-9.96729952e-01, 9.12780821e-02, 6.55573605e+00],\n", + " [ 1.28059260e-05, -4.64443066e-06, 1.00000000e+00]])\n", + "\n", + "x = 156.295617\n", + "y = 0\n", + "\n", + "print(np.atan2(mat[1, 0], mat[0, 0]) / math.pi * 180)\n", + "\n", + "print(cv2.perspectiveTransform(np.array([x, y]).reshape((1, 1, 2)), mat))\n", + "\n", + "nx, ny, na = autopilot.calc_position(np.linalg.inv(mat), x, y, math.pi / 2)\n", + "\n", + "print(nx, ny, na / math.pi * 180)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", - " [Pilot] translate: -1.5380446518826545 49.21369409309772\n", + " [Pilot] translate: -2.2243520368114386e-14 2.083887107525814e-15\n", " [Pilot] Drone Position: (47.97, -11.12)\n", " [Pilot] Angle: -11.3°\n", - " [Pilot] Target Index: 0\n", - " [Pilot] Target Position: [0. 0.]\n", - " [Pilot] Distance: 0.0\n", - "0.004129499336505804 50.0\n", + " [Pilot] Target Index: 1\n", + " [Pilot] Target Position: [347.32359333 -69.46471867]\n", + " [Pilot] Distance: 304.9907527982671\n", + "0.004129499336505887 50.0\n", + "47.96512465732409 -11.122053958129657\n", "\n", " [Pilot] translate: -0.773373251191613 49.85522586610471\n", " [Pilot] Drone Position: (96.75, -21.44)\n", @@ -69,7 +104,8 @@ " [Pilot] Target Index: 1\n", " [Pilot] Target Position: [347.32359333 -69.46471867]\n", " [Pilot] Distance: 304.9907527982671\n", - "0.003525142621744054 50.0\n", + "0.003525142621744165 50.0\n", + "96.7474713468631 -21.438269533988773\n", "\n", " [Pilot] translate: -0.4238729464016157 49.716078726704744\n", " [Pilot] Drone Position: (145.50, -31.21)\n", @@ -77,7 +113,8 @@ " [Pilot] Target Index: 1\n", " [Pilot] Target Position: [347.32359333 -69.46471867]\n", " [Pilot] Distance: 255.13708614259195\n", - "0.0019957432639431116 50.0\n", + "0.0019957432639432504 50.0\n", + "145.49554314093962 -31.210354705238323\n", "\n", " [Pilot] translate: -0.5072011337859063 49.476309464169745\n", " [Pilot] Drone Position: (194.03, -40.83)\n", @@ -85,7 +122,8 @@ " [Pilot] Target Index: 1\n", " [Pilot] Target Position: [347.32359333 -69.46471867]\n", " [Pilot] Distance: 205.42141613269357\n", - "0.0008527889497844943 50.0\n", + "0.0008527889497846608 50.0\n", + "194.02953935562192 -40.83395396007064\n", "\n", " [Pilot] translate: 0.11201474583194077 50.52107985317721\n", " [Pilot] Drone Position: (243.70, -50.06)\n", @@ -93,7 +131,8 @@ " [Pilot] Target Index: 1\n", " [Pilot] Target Position: [347.32359333 -69.46471867]\n", " [Pilot] Distance: 155.9448225243591\n", - "0.0006948810462804877 50.0\n", + "0.0006948810462807098 50.0\n", + "243.70149729872674 -50.058304936308375\n", "\n", " [Pilot] translate: 0.0398933070509686 50.472489376060885\n", " [Pilot] Drone Position: (293.34, -59.18)\n", @@ -102,65 +141,154 @@ " [Pilot] Target Position: [347.32359333 -69.46471867]\n", " [Pilot] Distance: 105.42365806674246\n", "-0.005757684837339738 50.0\n", + "293.3429112368944 -59.17991649962423\n", "\n", " [Pilot] translate: -0.5779730654942618 49.87019784134602\n", " [Pilot] Drone Position: (342.21, -69.12)\n", " [Pilot] Angle: -10.8°\n", " [Pilot] Target Index: 1\n", " [Pilot] Target Position: [347.32359333 -69.46471867]\n", - " [Pilot] Distance: 54.951716943628604\n", + " [Pilot] Distance: 54.95171694362866\n", "-0.6 50.0\n", + "342.21488384208834 -69.12481364723727\n", "\n", " [Pilot] translate: -0.3671006590822984 49.56271477490151\n", " [Pilot] Drone Position: (376.83, -104.60)\n", " [Pilot] Angle: -45.3°\n", " [Pilot] Target Index: 2\n", " [Pilot] Target Position: [ 295.22505433 -347.32359333]\n", - " [Pilot] Distance: 282.1393363134888\n", + " [Pilot] Distance: 282.13933631348874\n", "-0.6 50.0\n", + " [Pilot] translate: 152.48531037725778 -11.363336940584135\n", + "376.82813934296445 -104.60043911576437\n", "\n", " [Pilot] translate: 0.814116178363097 49.83161557706345\n", " [Pilot] Drone Position: (386.60, -153.47)\n", " [Pilot] Angle: -79.6°\n", " [Pilot] Target Index: 2\n", " [Pilot] Target Position: [ 295.22505433 -347.32359333]\n", - " [Pilot] Distance: 256.07341344997275\n", + " [Pilot] Distance: 256.0734134499727\n", "-0.6 50.0\n", - "\n" - ] - }, - { - "ename": "IndexError", - "evalue": "list index out of range", - "output_type": "error", - "traceback": [ - "\u001b[31m---------------------------------------------------------------------------\u001b[39m", - "\u001b[31mIndexError\u001b[39m Traceback (most recent call last)", - "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 3\u001b[39m\n\u001b[32m 1\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[32m1\u001b[39m, \u001b[32m20\u001b[39m):\n\u001b[32m 2\u001b[39m \u001b[38;5;28mprint\u001b[39m()\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m command = autopilot.handle(\u001b[43mimgs\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m)\n\u001b[32m 4\u001b[39m \u001b[38;5;28mprint\u001b[39m(command.dangle, command.velocity)\n\u001b[32m 5\u001b[39m vzm.set_target_index(autopilot.target_idx)\n", - "\u001b[31mIndexError\u001b[39m: list index out of range" + " [Pilot] translate: 223.2717832052611 -50.70818702916637\n", + "386.60059278853873 -153.47120593833623\n" ] } ], "source": [ - "for i in range(1, 20):\n", + "for i in range(1, 10):\n", " print()\n", " command = autopilot.handle(imgs[i])\n", " print(command.dangle, command.velocity)\n", " vzm.set_target_index(autopilot.target_idx)\n", " vzm.update_drone_trajectory(autopilot.x, autopilot.y)\n", " pos = autopilot.get_position_by_chunk()\n", - " if pos is not None:\n", - " print(pos)\n", - " vzm.update_global_map(pos[0], pos[1])\n", + " print(autopilot.x, autopilot.y)\n", + " vzm.update_global_map(0, 0)\n", " vzm.pause(2.5)\n" ] }, + { + "cell_type": "code", + "execution_count": 96, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[[ 0.18926711 -21.79121458]]]\n", + "0.40152777578126847\n" + ] + } + ], + "source": [ + "mat = np.array([[ 9.18211825e-01, -3.88263543e-01, 4.35846954e-01],\n", + " [ 3.89868385e-01, 9.19587423e-01, -2.49399645e+01],\n", + " [-3.56258907e-06, 3.01183248e-06, 1.00000000e+00]])\n", + "\n", + "pos = np.array([1., 3.]).reshape((1, 1, 2))\n", + "print(cv2.perspectiveTransform(pos, mat))\n", + "\n", + "(mat) @ np.array([1., 3., 1.]).reshape((3))\n", + "print(np.atan2(mat[1, 0], mat[0, 0]))" + ] + }, + { + "cell_type": "code", + "execution_count": 95, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[[-0.70710678 2.12132034]\n", + " [ 0.70710678 3.53553391]\n", + " [-2.12132034 3.53553391]\n", + " [ 0. 0. ]]\n", + "[[ 7.07106671e-01 -7.07106696e-01 -9.13243653e-17]\n", + " [ 7.07106528e-01 7.07106788e-01 3.65297461e-16]\n", + " [-7.16496367e-08 2.10734232e-09 1.00000000e+00]]\n", + "0.7853980622449934\n" + ] + } + ], + "source": [ + "ang = np.radians(45)\n", + "rot = np.array([\n", + " [np.cos(ang), -np.sin(ang)],\n", + " [np.sin(ang), np.cos(ang)]\n", + "])\n", + "\n", + "pts = np.array([\n", + " [1., 2.],\n", + " [3., 2.],\n", + " [1., 4.],\n", + " [0., 0.]\n", + "])\n", + "\n", + "dst = pts @ rot.T\n", + "\n", + "mat, mask = cv2.findHomography(pts, pts @ rot.T)\n", + "# print(np.atan2(mat[0, 1], mat[0, 0]))\n", + "# pts @ rot.T\n", + "pts @ mat[0:2]\n", + "# print(cv2.perspectiveTransform(pts[0].reshape((-1, 1, 2)), mat))\n", + "print(dst)\n", + "print(mat)\n", + "\n", + "\n", + "mat[0:2, 0:2] @ pts.T\n", + "\n", + "print(np.atan2(mat[1, 0], mat[0, 0]))\n" + ] + }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": 98, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(1, 2, 3)\n" + ] + } + ], + "source": [ + "def a(*args):\n", + " print(args)\n", + "\n", + "a(*(1, 2, 3))" + ] } ], "metadata": { diff --git a/timer.py b/timer.py new file mode 100644 index 0000000..b3fe31b --- /dev/null +++ b/timer.py @@ -0,0 +1,31 @@ +import time + +class Timer: + elapsed: float + last_enabled: float + enabled: bool + + def __init__(self): + self.elapsed = 0. + self.enabled = False + + def get_diff(self) -> float: + if not self.enabled: return 0. + return time.time() - self.last_enabled + + def start(self): + if self.enabled: return + self.enabled = True + self.last_enabled = time.time() + + def stop(self): + self.elapsed += self.get_diff() + self.enabled = False + + def get_elapsed(self) -> float: + return self.elapsed + self.get_diff() + + def reset(self): + self.elapsed = 0. + self.enabled = False + self.last_enabled = 0.