Write lots of code

undistort.py

@@ -0,0 +1,160 @@
+import os
+import sys
+
+import cv2
+import numpy as np
+import tkinter as tk
+from tkinter import filedialog, messagebox
+from PIL import Image, ImageTk
+from scipy.optimize import minimize_scalar
+
+from src import videoCrop
+from src.imageTools import radialUndistort, perspectiveUndistort, getFrameFromVideo
+
+class AdvancedDistortionCorrectionTool:
+    def __init__(self, master):
+        self.master = master
+        self.master.title("Advanced Distortion Correction Tool")
+        self.master.geometry("1200x800")
+
+        self.image = None
+        self.original_image = None
+        self.displayed_image = None
+        self.k = 0  # Initialize K value
+        self.dots = []
+        self.dragging = None
+
+        # Create UI elements
+        # self.load_button = tk.Button(self.master, text="Load Image", command=self.load_image)
+        # self.load_button.pack(pady=10)
+
+        self.canvas = tk.Canvas(self.master)
+        self.canvas.pack()
+        # self.canvas.width
+        self.canvas.bind("<Button-1>", self.on_click)
+        self.canvas.bind("<B1-Motion>", self.on_drag)
+        self.canvas.bind("<ButtonRelease-1>", self.on_release)
+
+        self.k_label = tk.Label(self.master, text="K value: 0")
+        self.k_label.pack(pady=5)
+
+        self.k_slider = tk.Scale(self.master, from_=-1.0, to=1.0, resolution=0.01, orient=tk.HORIZONTAL, length=300, command=self.update_k)
+        self.k_slider.pack(pady=10)
+
+        self.zoom_slider = tk.Scale(self.master, from_=0.5, to=5.0, resolution=0.01, orient=tk.HORIZONTAL, label="Zoom", length=300, command=self.update_zoom)
+        self.zoom_slider.pack(pady=10)
+
+        self.perspective_var = tk.BooleanVar()
+        self.perspective_check = tk.Checkbutton(self.master, text="Apply Perspective Correction", variable=self.perspective_var, command=self.update_persp)
+        self.perspective_check.pack(pady=10)
+
+        self.save_button = tk.Button(self.master, text="Save Image", command=self.save_image, state=tk.DISABLED)
+        self.save_button.pack(pady=10)
+
+        self.display_image_scale = 0.5
+
+        self.k_value = 0
+        self.zoom_value = 1
+
+    def update_zoom(self, zoom_value):
+        self.zoom_value = zoom_value
+        self.update_image(self.k_value, zoom_value)
+
+    def update_k(self, k_value):
+        self.k_value = k_value
+        self.update_image(k_value, self.zoom_value)
+
+    def update_persp(self):
+        self.update_image(self.k_value, self.zoom_value)
+
+    def load_image(self, image):
+        print(image)
+        self.original_image = image
+        self.save_button.config(state=tk.NORMAL)
+        height, width = self.original_image.shape[:2]
+        self.dots = [
+            [10, 10],
+            [width-10, 10],
+            [width-10, height-10],
+            [10, height-10]
+        ]
+        self.update_image(self.k_value, self.zoom_value)
+
+    def display_image(self):
+        if self.image is not None:
+            self.displayed_image = self.image.copy()
+            if not self.perspective_var.get():
+                for x, y in self.dots:
+                    cv2.circle(self.displayed_image, (int(x), int(y)), 15, (0, 255, 0), -1)
+            image = cv2.cvtColor(self.displayed_image, cv2.COLOR_BGR2RGB)
+            image = Image.fromarray(image)
+            image.thumbnail((image.width*self.display_image_scale, image.height*self.display_image_scale), Image.LANCZOS)
+            photo = ImageTk.PhotoImage(image)
+            self.canvas.config(width=photo.width(), height=photo.height())
+            self.canvas.create_image(0, 0, anchor=tk.NW, image=photo)
+            self.canvas.image = photo
+
+    def update_image(self, k_value, zoom_value):
+        if self.original_image is None:
+            return
+
+        self.image = radialUndistort(self.original_image, float(k_value), float(zoom_value))
+
+        if self.perspective_var.get():
+            self.image = perspectiveUndistort(self.image, self.dots)
+
+        self.display_image()
+        self.k_label.config(text=f"K value: {self.k:.2f}")
+
+    def count_lines(self, image):
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
+        lines = cv2.HoughLines(edges, 1, np.pi/180, 200)
+        return len(lines) if lines is not None else 0
+
+    def on_click(self, event):
+        x, y = event.x / self.display_image_scale, event.y / self.display_image_scale
+        for i, (dot_x, dot_y) in enumerate(self.dots):
+            if abs(x - dot_x) < 50 and abs(y - dot_y) < 50:
+                self.dragging = i
+                break
+
+
+    def on_drag(self, event):
+        if self.dragging is not None:
+            self.dots[self.dragging] = [event.x / self.display_image_scale, event.y / self.display_image_scale]
+            self.update_image(self.k_value, self.zoom_value)
+
+    def on_release(self, event):
+        self.dragging = None
+
+    def save_image(self):
+        print(",".join([
+            str(self.k_value),
+            str(self.zoom_value),
+
+            str(self.dots[0][0]),
+            str(self.dots[0][1]),
+            str(self.dots[1][0]),
+            str(self.dots[1][1]),
+            str(self.dots[2][0]),
+            str(self.dots[2][1]),
+            str(self.dots[3][0]),
+            str( self.dots[3][1])
+        ]))
+
+if __name__ == "__main__":
+    if len(sys.argv) < 1:
+        print("Usage: undistort.py <event code>")
+    filename = sys.argv[1] + "/" + (os.listdir(sys.argv[1])[0])
+
+    root = tk.Tk()
+    app = AdvancedDistortionCorrectionTool(root)
+
+    start, stop = videoCrop.videoCrop(filename)
+    if start != stop:
+        print("Found video start and end")
+
+    app.load_image(getFrameFromVideo(filename, start))
+
+    root.mainloop()