line_draw.py

from os import path
from sys import argv, exit

import cv2
import numpy as np

if len(argv) != 2:
    print("Usage: python line_draw.py path/to/image.png")
    exit()

image_path = argv[1]

image = cv2.imread(image_path)

shape_counter = 1
corner_counter = 1
corners = []
lines = []


def draw_lines(corners):
    if len(corners) >= 2:
        points = np.array(corners, dtype=np.int32)
        cv2.polylines(image, [points], isClosed=True, color=(0, 0, 255), thickness=2)


def save_image(image, filename):
    cv2.imwrite(filename, image)


def get_next_filename(index):
    return f"diff_images/approved_image{index}.png"


def mouse_callback(event, x, y, flags, param):
    global shape_counter, corner_counter, corners, lines

    if event == cv2.EVENT_LBUTTONDOWN:
        cv2.circle(image, (x, y), 5, (255, 255, 255), -1)
        cv2.imshow("Image", image)

        corner_name = f"shp{shape_counter}_x{(corner_counter - 1) % 4 + 1}, y{(corner_counter - 1) % 4 + 1}"
        print(f"{corner_name}: ({x}, {y})")

        corners.append((x, y))

        if corner_counter % 4 == 0:
            shape_counter += 1
            draw_lines(corners)
            lines.extend(corners)
            corners = []

        corner_counter += 1


index = 1

cv2.namedWindow("Image", cv2.WINDOW_NORMAL)
cv2.setMouseCallback("Image", mouse_callback)

while True:
    cv2.imshow("Image", image)

    choice = cv2.waitKey(0) & 0xFF
    if choice == ord("q"):
        break
    elif choice == ord("y") or choice == ord("Y"):
        filename = get_next_filename(index)
        while path.exists(filename):
            index += 1
            filename = get_next_filename(index)
        save_image(image, filename)
        index += 1

cv2.destroyAllWindows()