77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
import cv2
|
|
import cv2.aruco as aruco
|
|
# import dronekit
|
|
# from dronekit import connect
|
|
# from dronekit import VehicleMode
|
|
# from dronekit import mavutil
|
|
import logging
|
|
import time
|
|
|
|
logging.basicConfig(level='INFO')
|
|
|
|
class VisionSystem:
|
|
def __init__(self) -> None:
|
|
self.capture = cv2.VideoCapture(0)
|
|
self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
|
self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
|
_, frame = self.capture.read()
|
|
logging.info("Initialised camera")
|
|
logging.info('Resolution: ' + str(frame.shape[0]) + ' x ' + str(frame.shape[1]))
|
|
|
|
def find_marker(self, img, markerSize = 6, totalMarkers=250):
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
key = getattr(aruco, f'DICT_{markerSize}X{markerSize}_{totalMarkers}')
|
|
aruco_dict = aruco.Dictionary_get(key)
|
|
aruco_param = aruco.DetectorParameters_create()
|
|
bboxs, ids, rejected = aruco.detectMarkers(gray, aruco_dict, parameters=aruco_param)
|
|
if ids:
|
|
tl = bboxs[0][0][0][0], bboxs[0][0][0][1]
|
|
tr = bboxs[0][0][1][0], bboxs[0][0][1][1]
|
|
br = bboxs[0][0][2][0], bboxs[0][0][2][1]
|
|
bl = bboxs[0][0][3][0], bboxs[0][0][3][1]
|
|
# logging.info(f"{tl} | {tr}")
|
|
# logging.info(f"{bl} | {br}")
|
|
# logging.info(f"-----------")
|
|
|
|
# Get centre point of marker
|
|
total_x = [p[0] for p in [tr, tl, bl, br]]
|
|
total_y = [p[1] for p in [tr, tl, bl, br]]
|
|
centroid = (sum(total_x) / 4, sum(total_y) / 4)
|
|
|
|
# Get offset from centre
|
|
offset_y = (centroid[0] - (1280/2)) / 1280
|
|
offset_x = (centroid[1] - (720/2)) / 720
|
|
clean_offset_y = float("{0:.4f}".format(offset_y))
|
|
clean_offset_x = float("{0:.4f}".format(offset_x))
|
|
|
|
return (clean_offset_x, clean_offset_y)
|
|
return None
|
|
|
|
def calculate_angular_offset(self, pixel_offset):
|
|
'''
|
|
-0.5 -> 0.5
|
|
'''
|
|
logging.info("Calculating offset PWM signal")
|
|
x = pixel_offset[0]
|
|
y = float("{0:.1f}".format((pixel_offset[1] + 1.5) * 1000))
|
|
logging.info(f"Calculated offset: {y}")
|
|
return (x, y)
|
|
|
|
def get_pallet_angle(self):
|
|
while True:
|
|
success, img = self.capture.read()
|
|
logging.warning(f"Capture success: {success}")
|
|
logging.info("Looking for marker...")
|
|
raw_cargo_offset = self.find_marker(img)
|
|
logging.info(f"Data from find_marker(img): {raw_cargo_offset}")
|
|
# cv2.imshow('img',img)
|
|
if raw_cargo_offset:
|
|
logging.info(f"Found marker in get_pallet_angle {raw_cargo_offset}")
|
|
# cargo_offset = self.calculate_angular_offset(raw_cargo_offset)
|
|
# logging.info(f"PWM signal: {cargo_offset[1]}")
|
|
# return cargo_offset
|
|
# return None
|
|
|
|
v = VisionSystem()
|
|
v.get_pallet_angle()
|