common.py

import argparse
import logging

import numpy as np
import os

# Default file names.
FILE_3D = "3D.txt"
FILE_2D = "2D.txt"

# The default logging level.
LEVEL = "ERROR"


def logs(
        level: str = LEVEL
) -> None:
    """
    Configure logging.
    :param level: The level to set.
    :return: Nothing.
    """
    # Configure the logging.
    level = level.upper()
    if level == "CRITICAL" or level == "FATAL":
        level = logging.CRITICAL
    elif level == "ERROR":
        level = logging.ERROR
    elif level == "WARNING" or level == "WARN":
        level = logging.WARNING
    elif level == "INFO":
        level = logging.INFO
    elif level == "DEBUG":
        level = logging.DEBUG
    else:
        level = logging.NOTSET
    logging.basicConfig(level=level, format="%(asctime)s | %(levelname)s | %(message)s")


def load_file(
        file: str,
        dimensions: int = 2
) -> list[np.ndarray]:
    """
    Load a file of coordinates where points of the given dimensions are on each line separated by spaces.
    :param file: The path to the coordinate file.
    :param dimensions: The number of dimensions for each line of the coordinates.
    :return: A list of numpy arrays, each of which represents a data point with the given number of coordinates.
    """
    # Nothing to return if the file does not exist.
    if not os.path.exists(file):
        logging.error(f"File '{file}' does not exist.")
        return []
    # Ensure valid dimensions.
    if dimensions < 2:
        logging.warning(f"Dimensions of {dimensions} invalid, setting to 2 instead.")
        dimensions = 2
    # Read all points in the file.
    points = []
    with open(file, "r") as f:
        # Enumerate so we get line numbers in case we need to debug anything.
        for num, line in enumerate(f):
            try:
                # Split all elements of the line and convert them to floats.
                current = list(map(float, line.strip().split()))
                # Ensure this matches the number of dimensions we want to load.
                total = len(current)
                if total == dimensions:
                    # Append as a float32 numpy array to make it easily compatible with OpenCV.
                    points.append(np.array(current, dtype=np.float32))
                    continue
                # If this was the first line causing an error, assume it was indicating the total number of points.
                if num == 0 and total == 1:
                    logging.info(f"'{file}' is stated to contain {int(current[0])} points.")
                    continue
                # Ignore blank lines, but otherwise log the error.
                if total != 0:
                    logging.error(f"'{file}' at line {num + 1} had {total} coordinates instead of {dimensions}.")
                continue
            except Exception as e:
                logging.error(f"Could not parse '{file}' at line {num + 1} with content '{line}': {e}")
    if not points:
        logging.error(f"No {dimensions}-dimensional coordinates found in '{file}'.")
    return points


def points_valid(
        points_3d: list[np.ndarray],
        points_2d: list[np.ndarray]
) -> bool:
    """
    Check if the points can be used for calibration.
    :param points_3d: The 3D points in the form of [x, y, z] for each point.
    :param points_2d: The 2D points in the form of [x, y] for each point.
    :return: True if they can, false otherwise.
    """
    # Ensure the number of points match and there must be at least six of them.
    total_3d = len(points_3d)
    return total_3d >= 6 and total_3d == len(points_2d) 


def get_size(
        points_2d: list[np.ndarray],
        width: int = 1,
        height: int = 1
) -> tuple[int, int]:
    """
    Get the minimum size of the image based on the pixel values.
    :param points_2d: The 2D points in the form of [x, y] for each point.
    :param width: The desired camera width which will be used unless a point exceeds it.
    :param height: The desired camera height which will be used unless a point exceeds it.
    :return: A tuple with the width and height.
    """
    # At minimum, we must have one pixel.
    width = max(1, width)
    height = max(1, height)
    # Check every point.
    for point in points_2d:
        # See if this exceeds the width.
        w = int(np.ceil(point[0]))
        if w > width:
            width = w
        # See if this exceeds the height.
        h = int(np.ceil(point[1]))
        if h > height:
            height = h
    return width, height


def is_planar(
        points_3d: list[np.ndarray]
) -> bool:
    """
    Check if a list of points is planar. Note this does not check if the list is valid (more than six points).
    :param points_3d: The 3D points in the form of [x, y, z] for each point.
    :return: True if the points are planar, false otherwise.
    """
    # If any point has a non-zero Z value, it is non-planar.
    for point in points_3d:
        if point[2] != 0:
            return False
    return True


def load_files(
        file_3d: str = FILE_3D,
        file_2d: str = FILE_2D
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """
    Load corresponding 2D and 3D files of coordinates where points in each are on each line separated by spaces.
    :param file_3d: The path to the 3D coordinates file.
    :param file_2d: The path to the 2D coordinates file.
    :return: A list of numpy arrays for each file with each entry being a data point.
    """
    # Load each file.
    points_3d = load_file(file_3d, 3)
    points_2d = load_file(file_2d)
    # Ensure the number of points match.
    if points_valid(points_3d, points_2d):
        return points_3d, points_2d
    logging.error(f"Cannot calibrate using the points in '{file_3d}' and '{file_2d}'.")
    return [], []


def display_points(
        points: list[np.ndarray],
        title: str or None = None
) -> None:
    """
    Display the points in a list.
    :param points: The points.
    :param title: The title to display with.
    :return:
    """
    # Format the title.
    if title is None:
        title = ""
    else:
        title += ": "
    total = len(points)
    # Nothing to do if there are points.
    if total == 0:
        print(f"{title}0 points")
    # Otherwise, print all points.
    s = f"{title}{total} points\n{points[0]}"
    for i in range(1, total):
        s += f"\n{points[i]}"
    print(s)


def parse(
        description: str = "Computer Vision Program"
) -> tuple[str, str]:
    """
    Parse the command line arguments.
    :param description: The description to display.
    :return: The arguments for the paths to the 3D and 2D files.
    """
    # Add an argument parser.
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("-p", "--three", type=str, default=FILE_3D, help="The 3D points file.")
    parser.add_argument("-i", "--two", type=str, default=FILE_2D, help="The 2D points file.")
    parser.add_argument("-l", "--level", type=str, default=LEVEL, help="The logging level.")
    args = parser.parse_args()
    # Configure the logging.
    logs(args.level)
    return args.three, args.two


def main_common(
        file_3d: str = FILE_3D,
        file_2d: str = FILE_2D
) -> None:
    """
    Handle the main testing of the shared logic in this file.
    :param file_3d: The path to the 3D coordinates file.
    :param file_2d: The path to the 2D coordinates file.
    :return: Nothing.
    """
    # Load each of the files.
    points_3d, points_2d = load_files(file_3d, file_2d)
    # Display the loaded points.
    display_points(points_3d, "3D")
    display_points(points_2d, "2D")


if __name__ == "__main__":
    three, two = parse("Camera Calibration Helper Methods")
    main_common(three, two)