xhs-note-crawling/workplace/draw-point.py

from __future__ import annotations

import sys
from datetime import datetime
from pathlib import Path
from typing import Iterable, Sequence

import cv2

_REPOSITORY_ROOT_DIRECTORY = Path(__file__).resolve().parent.parent
_FULL_SCREEN_SCREENSHOT_PNG_FILE_PATH = _REPOSITORY_ROOT_DIRECTORY / "output" / "screenshot.png"
_POINT_DRAWING_OUTPUT_PARENT_DIRECTORY_PATH = _REPOSITORY_ROOT_DIRECTORY / "output"

_GREEN_POINT_COLOR_BGR_TUPLE = (0, 255, 0)
_POINT_CIRCLE_RADIUS_PIXELS_INTEGER = 6
_POINT_CIRCLE_OUTLINE_THICKNESS_PIXELS_INTEGER = -1


def draw_green_points_on_image_file_and_save_to_output_path(
    points_xy_coordinate_sequence: Sequence[Sequence[float]],
    input_image_file_path: str | Path,
    output_image_file_path: str | Path,
) -> Path:
    input_image_file_path_resolved = Path(input_image_file_path).resolve()
    output_image_file_path_resolved = Path(output_image_file_path).resolve()
    output_image_file_path_resolved.parent.mkdir(parents=True, exist_ok=True)
    input_image_bgr_numpy_array = cv2.imread(
        str(input_image_file_path_resolved),
        cv2.IMREAD_COLOR,
    )
    output_image_bgr_numpy_array = input_image_bgr_numpy_array.copy()
    points_iterable: Iterable[Sequence[float]] = points_xy_coordinate_sequence
    for point_xy_pair in points_iterable:
        center_x_integer = int(round(float(point_xy_pair[0])))
        center_y_integer = int(round(float(point_xy_pair[1])))
        cv2.circle(
            output_image_bgr_numpy_array,
            (center_x_integer, center_y_integer),
            _POINT_CIRCLE_RADIUS_PIXELS_INTEGER,
            _GREEN_POINT_COLOR_BGR_TUPLE,
            _POINT_CIRCLE_OUTLINE_THICKNESS_PIXELS_INTEGER,
        )
    cv2.imwrite(str(output_image_file_path_resolved), output_image_bgr_numpy_array)
    return output_image_file_path_resolved


def draw_point_xy_on_current_screenshot_png_and_save_timestamped_png_file(
    points_xy_coordinate_sequence: Sequence[Sequence[float]],
) -> Path:
    timestamp_string_for_output_filename = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    keypoints_preview_png_file_path = (
        _POINT_DRAWING_OUTPUT_PARENT_DIRECTORY_PATH
        / f"keypoints-{timestamp_string_for_output_filename}.png"
    )
    return draw_green_points_on_image_file_and_save_to_output_path(
        points_xy_coordinate_sequence,
        _FULL_SCREEN_SCREENSHOT_PNG_FILE_PATH,
        keypoints_preview_png_file_path,
    )


def start(command_line_argument_string_list: list[str] | None = None) -> int:
    command_line_arguments = (
        command_line_argument_string_list
        if command_line_argument_string_list is not None
        else sys.argv[1:]
    )
    input_image_file_path_string = command_line_arguments[0].strip()
    output_image_file_path_string = command_line_arguments[1].strip()
    coordinate_float_string_list = command_line_arguments[2:]
    points_xy_coordinate_pair_list: list[list[float]] = []
    pair_index_integer = 0
    while pair_index_integer + 1 < len(coordinate_float_string_list):
        x_coordinate_float = float(coordinate_float_string_list[pair_index_integer].strip())
        y_coordinate_float = float(
            coordinate_float_string_list[pair_index_integer + 1].strip()
        )
        points_xy_coordinate_pair_list.append([x_coordinate_float, y_coordinate_float])
        pair_index_integer += 2
    draw_green_points_on_image_file_and_save_to_output_path(
        points_xy_coordinate_pair_list,
        input_image_file_path_string,
        output_image_file_path_string,
    )
    return 0


def main() -> int:
    return start()


__all__ = [
    "draw_green_points_on_image_file_and_save_to_output_path",
    "draw_point_xy_on_current_screenshot_png_and_save_timestamped_png_file",
    "start",
    "main",
]


if __name__ == "__main__":
    raise SystemExit(main())