#%%
import cv2
import numpy as np
import os

# Initialize global variables
points = []
drawing = False

# Function to clear the drawn points
def clear_points():
    global points
    points = []

# Mouse callback function
def draw_polygon(event, x, y, flags, param):
    global points, drawing, img

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        points.append((x, y))

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            img_copy = img.copy()
            for i in range(1, len(points)):
                cv2.line(img_copy, points[i - 1], points[i], (255, 0, 0), 2)
            if len(points) > 0:
                cv2.circle(img_copy, points[-1], 3, (0, 0, 255), -1, lineType=cv2.LINE_AA)  # Hiển thị điểm chọn của chuột
            cv2.imshow('image', img_copy)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        points.append((x, y))
        pts = np.array(points, np.int32)
        pts = pts.reshape((-1, 1, 2))
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        cv2.fillPoly(mask, [pts], 255)
        cv2.imwrite(mask_path, mask)
        cv2.imshow('image', img)

# Function to process images in a folder
def process_images_in_folder(folder_path):
    global img, mask_path

    for img_name in os.listdir(folder_path):
        if img_name.endswith('.jpg'):
            img_path = os.path.join(folder_path, img_name)
            mask_path = os.path.join(folder_path, f'{os.path.splitext(img_name)[0]}_mask.jpg')
            img = cv2.imread(img_path)

            # Create a window and bind the mouse callback function
            cv2.namedWindow('image')
            cv2.setMouseCallback('image', draw_polygon)

            while True:
                cv2.imshow('image', img)
                k = cv2.waitKey(1) & 0xFF
                if k == 27:  # Press 'ESC' to exit
                    break

            clear_points()
            cv2.destroyAllWindows()

# Define folders to process
folder_path=r'C:\Users\20240805\Documents\GitHub\AD-CLIP\data\4inlab\shinpyung\train\anomaly'
process_images_in_folder(folder_path)
# %%