Untitled

mail@pastecode.io avatar
unknown
python
a year ago
1.9 kB
2
Indexable
Never
# This code only uses the rotation matrix for vizualisation

import os
import numpy as np
import cv2

# Define the length of the arrows in pixels
arrow_length = 50

# Define the colors of the arrows
arrow_colors = [(0,0,255), (0,255,0), (255,0,0)] # BGR format

# Define the starting point for the arrows (top-right corner)
start_point = (50, 50)

# Path to the folder containing the RGB files
folder_path = r"C:\Users\KoenF\Documents\MP\EIGEN_CODE\data\vizualisationCode\bleach\rgb" #Zonder / op het einde

# Loop over all PNG files in the folder
for filename in os.listdir(folder_path):
    if filename.endswith(".png"):
        # Construct the paths for the pose file and the output file
        base_filename = os.path.splitext(filename)[0]
        pose_file = os.path.join(folder_path, "../poses/{}.txt".format(base_filename))
        output_file = os.path.join(folder_path, "../vizualisation/{}.png".format(base_filename))

        # Read in the pose file
        pose = np.loadtxt(pose_file)

        # Define the endpoints for the arrows based on the pose
        end_points = [
            (int(start_point[0] + arrow_length * pose[2][0]), int(start_point[1] - arrow_length * pose[2][1])),
            (int(start_point[0] + arrow_length * pose[0][0]), int(start_point[1] - arrow_length * pose[0][1])),
            (int(start_point[0] + arrow_length * pose[1][0]), int(start_point[1] - arrow_length * pose[1][1]))
        ]

        # Load the RGB image
        image_file = os.path.join(folder_path, filename)
        image = cv2.imread(image_file)

        # Draw the arrows on the image
        for end_point, color in zip(end_points, arrow_colors):
            image = cv2.arrowedLine(image, start_point, end_point, color, thickness=2)

        # Save the image with arrows
        cv2.imwrite(output_file, image)