draw_poly.py

import cv2
import tkinter as tk
from tkinter import messagebox
from PIL import Image, ImageTk
import json
from typing import List, Tuple
from shapely.geometry import Polygon, box
from shapely import intersects


class CurveDrawer:
    def __init__(self, img_entry, img_exit):
        # Store original images (BGR)
        self.orig_entry = img_entry
        self.orig_exit = img_exit

        # Native dimensions
        self.h_entry, self.w_entry = img_entry.shape[:2]
        self.h_exit, self.w_exit = img_exit.shape[:2]

        # UI
        self.root = tk.Tk()
        self.root.title("Draw Polylines: Entry (Left) | Exit (Right)")

        # Get screen size for smart display scaling
        temp = tk.Tk()
        temp.withdraw()
        screen_w = temp.winfo_screenwidth()
        screen_h = temp.winfo_screenheight()
        temp.destroy()

        max_h = min(screen_h - 250, 900)
        max_w_per_side = (screen_w // 2) - 60

        # --- Process Entry ---
        scale_entry = 1.0
        if self.w_entry > max_w_per_side or self.h_entry > max_h:
            scale_w = max_w_per_side / self.w_entry
            scale_h = max_h / self.h_entry
            scale_entry = min(scale_w, scale_h)
        self.disp_w_entry = int(self.w_entry * scale_entry)
        self.disp_h_entry = int(self.h_entry * scale_entry)
        self.scale_x_entry = self.w_entry / self.disp_w_entry
        self.scale_y_entry = self.h_entry / self.disp_h_entry

        rgb_entry = cv2.cvtColor(img_entry, cv2.COLOR_BGR2RGB)
        pil_entry = Image.fromarray(rgb_entry).resize((self.disp_w_entry, self.disp_h_entry), Image.LANCZOS)
        self.photo_entry = ImageTk.PhotoImage(pil_entry)

        # --- Process Exit ---
        scale_exit = 1.0
        if self.w_exit > max_w_per_side or self.h_exit > max_h:
            scale_w = max_w_per_side / self.w_exit
            scale_h = max_h / self.h_exit
            scale_exit = min(scale_w, scale_h)
        self.disp_w_exit = int(self.w_exit * scale_exit)
        self.disp_h_exit = int(self.h_exit * scale_exit)
        self.scale_x_exit = self.w_exit / self.disp_w_exit
        self.scale_y_exit = self.h_exit / self.disp_h_exit

        rgb_exit = cv2.cvtColor(img_exit, cv2.COLOR_BGR2RGB)
        pil_exit = Image.fromarray(rgb_exit).resize((self.disp_w_exit, self.disp_h_exit), Image.LANCZOS)
        self.photo_exit = ImageTk.PhotoImage(pil_exit)

        # Polylines in NATIVE coordinates
        self.polyline_entry: List[Tuple[int, int]] = []
        self.polyline_exit: List[Tuple[int, int]] = []

        # Zone choices
        self.zone_entry = tk.StringVar(value="above")
        self.zone_exit = tk.StringVar(value="above")

        # === Entry ===
        frame_entry = tk.Frame(self.root)
        frame_entry.pack(side=tk.LEFT, padx=10, pady=10)
        self.canvas_entry = tk.Canvas(frame_entry, width=self.disp_w_entry, height=self.disp_h_entry)
        self.canvas_entry.pack()
        self.canvas_entry.create_image(0, 0, anchor=tk.NW, image=self.photo_entry)
        self.canvas_entry.bind("<Button-1>", self._on_click_entry)

        radio_entry = tk.Frame(frame_entry)
        radio_entry.pack(pady=5)
        tk.Label(radio_entry, text="Entry zone cut off:").pack(anchor=tk.W)
        tk.Radiobutton(radio_entry, text="Below line", variable=self.zone_entry, value="below").pack(anchor=tk.W)
        tk.Radiobutton(radio_entry, text="Above line", variable=self.zone_entry, value="above").pack(anchor=tk.W)

        # === Exit ===
        frame_exit = tk.Frame(self.root)
        frame_exit.pack(side=tk.RIGHT, padx=10, pady=10)
        self.canvas_exit = tk.Canvas(frame_exit, width=self.disp_w_exit, height=self.disp_h_exit)
        self.canvas_exit.pack()
        self.canvas_exit.create_image(0, 0, anchor=tk.NW, image=self.photo_exit)
        self.canvas_exit.bind("<Button-1>", self._on_click_exit)

        radio_exit = tk.Frame(frame_exit)
        radio_exit.pack(pady=5)
        tk.Label(radio_exit, text="Exit zone cut off:").pack(anchor=tk.W)
        tk.Radiobutton(radio_exit, text="Below line", variable=self.zone_exit, value="below").pack(anchor=tk.W)
        tk.Radiobutton(radio_exit, text="Above line", variable=self.zone_exit, value="above").pack(anchor=tk.W)

        # Buttons
        btn_frame = tk.Frame(self.root)
        btn_frame.pack(pady=10)
        tk.Button(btn_frame, text="OK", command=self._on_ok, width=10).grid(row=0, padx=5)
        tk.Button(btn_frame, text="Clear Entry", command=self._clear_entry, width=12).grid(row=1, padx=5)
        tk.Button(btn_frame, text="Clear Exit", command=self._clear_exit, width=12).grid(row=2, padx=5)

        self.lines_entry = []
        self.lines_exit = []

    def _on_click_entry(self, event):
        x_native = int(round(event.x * self.scale_x_entry))
        y_native = int(round(event.y * self.scale_y_entry))
        x_native = max(0, min(self.w_entry - 1, x_native))
        y_native = max(0, min(self.h_entry - 1, y_native))
        self.polyline_entry.append((x_native, y_native))
        self._redraw_entry()

    def _on_click_exit(self, event):
        x_native = int(round(event.x * self.scale_x_exit))
        y_native = int(round(event.y * self.scale_y_exit))
        x_native = max(0, min(self.w_exit - 1, x_native))
        y_native = max(0, min(self.h_exit - 1, y_native))
        self.polyline_exit.append((x_native, y_native))
        self._redraw_exit()

    def _redraw_entry(self):
        for lid in self.lines_entry:
            self.canvas_entry.delete(lid)
        self.lines_entry.clear()
        self.canvas_entry.create_image(0, 0, anchor=tk.NW, image=self.photo_entry)
        self._draw_polyline(self.canvas_entry, self.polyline_entry, self.scale_x_entry, self.scale_y_entry)

    def _redraw_exit(self):
        for lid in self.lines_exit:
            self.canvas_exit.delete(lid)
        self.lines_exit.clear()
        self.canvas_exit.create_image(0, 0, anchor=tk.NW, image=self.photo_exit)
        self._draw_polyline(self.canvas_exit, self.polyline_exit, self.scale_x_exit, self.scale_y_exit)

    def _draw_polyline(self, canvas, polyline_native, scale_x, scale_y):
        if len(polyline_native) == 0:
            return
        disp_pts = [
            (int(round(x / scale_x)), int(round(y / scale_y)))
            for (x, y) in polyline_native
        ]
        if len(disp_pts) == 1:
            x, y = disp_pts[0]
            canvas.create_oval(x-3, y-3, x+3, y+3, fill="red")
        for i in range(1, len(disp_pts)):
            x1, y1 = disp_pts[i-1]
            x2, y2 = disp_pts[i]
            canvas.create_line(x1, y1, x2, y2, fill="red", width=2)
            canvas.create_oval(x2-3, y2-3, x2+3, y2+3, fill="red")

    def _clear_entry(self):
        self.polyline_entry.clear()
        self._redraw_entry()

    def _clear_exit(self):
        self.polyline_exit.clear()
        self._redraw_exit()

    def _extend_to_nearest_vertical_edge(self, point, width):
        x, y = point
        if x < width / 2:
            return (0, y)
        else:
            return (width - 1, y)

    def _on_ok(self):
        if len(self.polyline_entry) < 2 or len(self.polyline_exit) < 2:
            messagebox.showwarning("Warning", "Both need ≥2 points.")
            return

        # Extend using native widths
        start_e = self._extend_to_nearest_vertical_edge(self.polyline_entry[0], self.w_entry)
        end_e = self._extend_to_nearest_vertical_edge(self.polyline_entry[-1], self.w_entry)
        final_entry = [start_e] + self.polyline_entry + [end_e]

        start_x = self._extend_to_nearest_vertical_edge(self.polyline_exit[0], self.w_exit)
        end_x = self._extend_to_nearest_vertical_edge(self.polyline_exit[-1], self.w_exit)
        final_exit = [start_x] + self.polyline_exit + [end_x]

        data = {
            "entry": {"polyline": final_entry, "zone": self.zone_entry.get(), "image_size": [self.w_entry, self.h_entry]},
            "exit": {"polyline": final_exit, "zone": self.zone_exit.get(), "image_size": [self.w_exit, self.h_exit]}
        }

        with open('polygons.json', 'w') as f:
            json.dump(data, f, indent=2)

        messagebox.showinfo("Success", "Saved successfully")
        self.root.quit()

    def show(self):
        self.root.mainloop()


# =============================================================================
# Updated loader: now uses stored image_size
# =============================================================================

def load_clipping_polygon(filepath: str, camera: str):
    """
    Returns (polygon, zone, (width, height)) for the given camera.
    No need to pass image size externally.
    """
    with open(filepath, 'r') as f:
        data = json.load(f)

    if camera not in data:
        raise KeyError(f"Camera '{camera}' not in {list(data.keys())}")

    item = data[camera]
    polyline = item["polyline"]
    zone = item["zone"]
    w, h = item["image_size"]

    if len(polyline) < 2:
        raise ValueError(f"Polyline for {camera} too short")

    if zone == "below":
        closed = polyline + [(w - 1, h - 1), (0, h - 1)]
    elif zone == "above":
        closed = polyline + [(w - 1, 0), (0, 0)]
    else:
        raise ValueError("Invalid zone")

    return Polygon(closed), zone


def rectangle_intersects_clipping_zone(rect: Tuple[Tuple[int, int], Tuple[int, int]],
                                       polygon: Polygon) -> bool:
    (x1, y1), (x2, y2) = rect
    rect_box = box(min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))
    return intersects(polygon, rect_box)


# =============================================================================
# Main
# =============================================================================

if __name__ == '__main__':
    entry_id = input('Entry camera id: ').strip()
    exit_id = input('Exit camera id: ').strip()
    cap0 = cv2.VideoCapture(int(entry_id) if entry_id else 0)
    cap1 = cv2.VideoCapture(int(exit_id) if exit_id else 1)
##    cap0 = cv2.VideoCapture("C:\\Users\\USER\\Downloads\\IMG_2907_2fps.mp4")
##    cap1 = cv2.VideoCapture(1)

    ret0, frame0 = cap0.read()
    ret1, frame1 = cap1.read()
    cap0.release()
    cap1.release()

    if not (ret0 and ret1):
        print("Failed to read videos")
        exit()

    drawer = CurveDrawer(frame0, frame1)
    drawer.show()

##if __name__ == '__main__':
##    def list_cameras(max_test=10):
##        cams = []
##        for i in range(max_test):
##            cap = cv2.VideoCapture(i)
##            if cap.isOpened():
##                ret, _ = cap.read()
##                if ret:
##                    cams.append(i)
##                cap.release()
##        return cams
##
####    print("Available cameras:", list_cameras())
##    entry_id = input('Entry camera id: ').strip()
##    exit_id = input('Exit camera id: ').strip()
####    cap0 = cv2.VideoCapture(int(entry_id) if entry_id else 0),
####    cap1 = cv2.VideoCapture(int(exit_id) if exit_id else 1)
##    cap0 = cv2.VideoCapture("C:\\Users\\USER\\Downloads\\IMG_2907_2fps.mp4")
##    cap1 = cv2.VideoCapture(1)
##    
##    
##    if not cap0.isOpened():
##        print(f"Error: Cannot open entry camera ({entry_id})")
##        exit()
##    if not cap1.isOpened():
##        print(f"Error: Cannot open exit camera ({exit_id})")
##        exit()
##
##    ret0, frame0 = cap0.read()
##    ret1, frame1 = cap1.read()
##    cap0.release()
##    cap1.release()
##    print(frame0.shape)
##    print(frame1.shape)
##    
##    if not (ret0 and ret1):
##        print("Failed to capture from one or both cameras")
##        exit()
##
##    # Resize to match if needed
##    h0, w0 = frame0.shape[:2]
##    h1, w1 = frame1.shape[:2]
##    if (h0, w0) != (h1, w1):
##        frame1 = cv2.resize(frame1, (w0, h0))
##
##    drawer = CurveDrawer(frame0, frame1)
##    drawer.show()