camera_test.py

import cv2
import cv2.aruco as aruco
import numpy as np

cap = cv2.VideoCapture("http://10.0.0.189:8080/videofeed")

# ArUco detector
aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_4X4_50)
detector = aruco.ArucoDetector(aruco_dict)

# Camera calibration
focal_length = 800
center = (320, 240)
camera_matrix = np.array([[focal_length, 0, center[0]],
                          [0, focal_length, center[1]],
                          [0, 0, 1]], dtype=float)
dist_coeffs = np.zeros((4, 1))
marker_length = 2.0  # cm


def estimate_marker_pose(corners, marker_size, camera_matrix, dist_coeffs):
    """Estimate marker pose using solvePnP"""
    obj_points = np.array([[-marker_size / 2, marker_size / 2, 0],
                           [marker_size / 2, marker_size / 2, 0],
                           [marker_size / 2, -marker_size / 2, 0],
                           [-marker_size / 2, -marker_size / 2, 0]], dtype=np.float32)

    success, rvec, tvec = cv2.solvePnP(obj_points, corners.astype(np.float32),
                                       camera_matrix, dist_coeffs)
    if success:
        return rvec, tvec
    return None, None


def rotation_vector_to_euler_angles(rvec):
    """Convert rotation vector to Euler angles (roll, pitch, yaw) in degrees"""
    # Convert rotation vector to rotation matrix
    R, _ = cv2.Rodrigues(rvec)

    # Calculate Euler angles from rotation matrix
    # Different conventions exist - using this common one:
    sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])

    singular = sy < 1e-6

    if not singular:
        x = np.arctan2(R[2, 1], R[2, 2])
        y = np.arctan2(-R[2, 0], sy)
        z = np.arctan2(R[1, 0], R[0, 0])
    else:
        x = np.arctan2(-R[1, 2], R[1, 1])
        y = np.arctan2(-R[2, 0], sy)
        z = 0

    # Convert to degrees
    roll = np.degrees(x)
    pitch = np.degrees(y)
    yaw = np.degrees(z)

    return roll, pitch, yaw


def order_points_clockwise(points):
    """Order 4 points in clockwise order starting from top-left"""
    points = np.array(points)
    center = np.mean(points, axis=0)
    angles = []
    for point in points:
        angle = np.arctan2(point[1] - center[1], point[0] - center[0])
        angles.append(angle)
    sorted_indices = np.argsort(angles)
    return points[sorted_indices]


def calculate_plane_orientation(marker_positions_3d):
    """Calculate the orientation of the plane formed by 4 markers"""
    if len(marker_positions_3d) != 4:
        return None, None, None

    points_3d = np.array(marker_positions_3d)

    # Fit a plane to the 4 points
    centroid = np.mean(points_3d, axis=0)
    points_centered = points_3d - centroid

    # Use SVD to find the normal vector
    U, S, Vt = np.linalg.svd(points_centered)
    normal = Vt[2]  # Last row of Vt is the normal

    # Make sure normal points toward camera (negative Z direction)
    if normal[2] > 0:
        normal = -normal

    # Calculate Euler angles from normal vector
    # For a plane, we can interpret the normal as defining orientation
    nx, ny, nz = normal

    # Calculate roll, pitch from normal vector
    pitch = np.degrees(np.arcsin(-ny))
    roll = np.degrees(np.arctan2(nx, nz))

    # Yaw requires more complex calculation - we'll use the first marker's orientation
    yaw = 0  # Will be calculated per marker

    return roll, pitch, normal


print("✅ 6-DOF Medical Robot Operating Area Tracker")
print("📍 Tracking Position (X,Y,Z) + Orientation (Roll,Pitch,Yaw)...")

while True:
    ret, frame = cap.read()
    if not ret:
        break

    frame = cv2.resize(frame, (640, 480))
    corners, ids, rejected = detector.detectMarkers(frame)

    if ids is not None and len(ids) >= 4:
        frame = aruco.drawDetectedMarkers(frame, corners, ids)

        marker_positions_2d = []
        marker_positions_3d = []
        marker_orientations = []
        marker_info = []

        for i in range(len(ids)):
            corner = corners[i][0]
            center_2d = np.mean(corner, axis=0)

            # Estimate 3D pose
            rvec, tvec = estimate_marker_pose(corner, marker_length, camera_matrix, dist_coeffs)

            if rvec is not None:
                x, y, z = tvec.flatten()
                marker_positions_2d.append(center_2d)
                marker_positions_3d.append((x, y, z))

                # Calculate orientation for this marker
                roll, pitch, yaw = rotation_vector_to_euler_angles(rvec)
                marker_orientations.append((roll, pitch, yaw))

                marker_info.append((ids[i][0], center_2d, (x, y, z), (roll, pitch, yaw)))

                # Display marker ID and basic orientation
                cv2.putText(frame, f"ID:{ids[i][0]}",
                            (int(center_2d[0]) - 15, int(center_2d[1]) - 15),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 2)
                cv2.putText(frame, f"ID:{ids[i][0]}",
                            (int(center_2d[0]) - 15, int(center_2d[1]) - 15),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)

        # Create proper rectangle if we have 4 markers
        if len(marker_positions_3d) == 4:
            ordered_points = order_points_clockwise(marker_positions_2d)
            centroid_2d = np.mean(ordered_points, axis=0)
            centroid_3d = np.mean(marker_positions_3d, axis=0)

            # Calculate plane orientation
            plane_roll, plane_pitch, plane_normal = calculate_plane_orientation(marker_positions_3d)

            # Calculate average orientation from all markers
            avg_roll = np.mean([o[0] for o in marker_orientations])
            avg_pitch = np.mean([o[1] for o in marker_orientations])
            avg_yaw = np.mean([o[2] for o in marker_orientations])

            # Calculate dimensions
            points_3d = np.array(marker_positions_3d)
            min_coords = np.min(points_3d, axis=0)
            max_coords = np.max(points_3d, axis=0)
            dimensions = max_coords - min_coords

            # Draw the properly ordered rectangle
            points_int = ordered_points.astype(np.int32)
            cv2.polylines(frame, [points_int], True, (0, 255, 0), 3)

            # Fill with semi-transparent color
            overlay = frame.copy()
            cv2.fillPoly(overlay, [points_int], (0, 255, 0))
            cv2.addWeighted(overlay, 0.15, frame, 0.85, 0, frame)

            # Draw centroid and coordinate axes
            cv2.circle(frame, (int(centroid_2d[0]), int(centroid_2d[1])), 8, (255, 0, 0), -1)
            cv2.circle(frame, (int(centroid_2d[0]), int(centroid_2d[1])), 12, (255, 0, 0), 2)

            # Draw normal vector visualization
            if plane_normal is not None:
                normal_end = centroid_2d + plane_normal[:2] * 50  # Project to 2D
                cv2.arrowedLine(frame,
                                (int(centroid_2d[0]), int(centroid_2d[1])),
                                (int(normal_end[0]), int(normal_end[1])),
                                (0, 0, 255), 3, tipLength=0.3)

            # Display comprehensive 6-DOF information
            info_lines = [
                "6-DOF OPERATING AREA TRACKER",
                f"POSITION: X:{centroid_3d[0]:6.1f}cm Y:{centroid_3d[1]:6.1f}cm Z:{centroid_3d[2]:6.1f}cm",
                f"ORIENTATION: R:{avg_roll:5.1f}° P:{avg_pitch:5.1f}° Y:{avg_yaw:5.1f}°",
                f"PLANE TILT: Roll:{plane_roll:5.1f}° Pitch:{plane_pitch:5.1f}°",
                f"DIMENSIONS: {dimensions[0]:.1f} x {dimensions[1]:.1f} x {dimensions[2]:.1f} cm",
                f"MARKERS: {len(ids)}/4 - 6-DOF Tracking Active"
            ]

            for j, line in enumerate(info_lines):
                cv2.putText(frame, line, (10, 30 + j * 22),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 3)
                cv2.putText(frame, line, (10, 30 + j * 22),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

            # Print detailed 6-DOF data to console
            print(f"📍 6-DOF: Pos({centroid_3d[0]:6.1f}, {centroid_3d[1]:6.1f}, {centroid_3d[2]:6.1f}) cm | "
                  f"Rot(R:{avg_roll:5.1f}°, P:{avg_pitch:5.1f}°, Y:{avg_yaw:5.1f}°) | "
                  f"Plane(R:{plane_roll:5.1f}°, P:{plane_pitch:5.1f}°)")

    else:
        markers_found = len(ids) if ids is not None else 0
        cv2.putText(frame, f"6-DOF Tracking: Need 4 markers, found {markers_found}", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

    cv2.imshow("6-DOF Medical Robot Operating Area Tracker", frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()