mpenn/scripts/Labeling.py

import customtkinter as Ctk
import os
from PIL import Image, ImageTk, ImageDraw
from scripts.SaveData import SaveData
from icons.icons import Icons
import scripts.get_sys_info as system_code



class Labeling(Ctk.CTkFrame):
    def __init__(self,master,img_paths, output_path, callback, window_information, **kwargs):
        super().__init__(master, **kwargs)
        system_code.load_json_file()
        self.icons = Icons(path=system_code.data_path)
        self.data_saver = SaveData()
        self.my_font = Ctk.CTkFont(family="Berlin Sans FB", size=22)
        self.image_btn_size = system_code.btn_img_size
        #callback
        self.callback = callback
        self.window_information = window_information

        # Load and resize image as before
        self.create_folder_raw_image = self.icons.get_image("CreateFolder")
        self.open_folder_raw_image = self.icons.get_image("OpenFolder")
        self.source_folder_raw_image = self.icons.get_image("SourceFolder")
        self.load_button_images()
        # for the image show_process
        self.image_scale_init()
        # variables for the labeling itself
        self.index = 0 # when this gets opened we always want to start a 0.
        self.img_paths = img_paths
        self.active_output_path = output_path
        self.output_path = output_path

        # bounding boxes
        self.labeling_boxes =  []
        self.data_mode_index = system_code.data_modes.index(system_code.data_mode)
        self.labeling_box_index = 0
        self.original_factor = None
        self.img_factor_x = None
        self.img_factor_y = None
        self.resolution = None
        self.labeled = False
        self.mpenn_data = self.data_saver.search_in_folder(self.active_output_path)
        self.save_index = self.data_saver.find_highest_image_number(self.mpenn_data[0])

        self.start()
    
    def enable_keybinding(self):
        """
        r: Remove the picture
        j: jump in Time
        l: Labeling / Resize Switch
        o: Open Folder
        c: Create Folder 
        s: Switch Folder
        Return: Save Image
        """
        self.master.bind("<r>", self.delete_current_rectangle, add="+")
        self.master.bind("<j>", self.skip_time, add="+")
        self.master.bind("<l>", self.set_data_mode_binding, add="+")
        self.master.bind("<o>", self.open_new_folder, add="+")
        self.master.bind("<c>", self.create_new_folder, add="+")
        self.master.bind("<s>", self.source_folder_dialog, add="+")
        self.master.bind("<Return>", self.save_and_load, add="+")

    def disable_keybinding(self):
        self.master.unbind("<r>")
        self.master.unbind("<j>")
        self.master.unbind("<l>")
        self.master.unbind("<o>")
        self.master.unbind("<c>")
        self.master.unbind("<s>")
        self.master.unbind("<Return>")

    def update_active_output_path(self, output_path):
        self.active_output_path = output_path
        self.mpenn_data = self.data_saver.search_in_folder(self.active_output_path)
        self.save_index = self.data_saver.find_highest_image_number(self.mpenn_data[0])
    
    def start(self):
        self.create_labeling()
        self.show_img(self.img_paths[self.index]) 

    def open_new_folder(self):
        """
        Callback for when a new folder is chosen.
        """
        self.callback(self.output_path, 1)  # Call the callback passing the current value

    def create_new_folder(self):
        """
        Callback for when a new folder is chosen.
        """
        self.callback(self.output_path, 0)  # Call the callback passing the current value
    
    def return_image_information(self):
        # Extract the directory of the image file
        image_directory = os.path.dirname(self.img_paths[self.index])
        
        # Now extract the folder name
        desired_folder = os.path.basename(image_directory)
        desired_name = os.path.basename(self.img_paths[self.index])
        
        # Creating a tuple with the extracted information
        text_information = (desired_name, desired_folder, self.active_output_path)
        
        # Presumed function to show or use the information
        self.window_information(text_information)

    
    def source_folder_dialog(self):
        self.callback(None, 2)
    
    def image_scale_init(self):
        self.original_image = None
        self.displayed_image = None
        self.dragging = False
        self.resize_pending = False
        self.after_id = None
        self.tk_image = None
        self.rect = None  # Keep track of the rectangle element
        self.start_x = None
        self.start_y = None
        self.end_x = None  # Track the end position of the drag
        self.end_y = None
        self.image_position = (0, 0)  # Initialize image position
        self.tk_cropped_image = None
        self.save_cropped = None


    # the Labeling part
    def create_labeling(self):
        """adapt button size to window_size"""
        self.big_canvas = Ctk.CTkCanvas(self,background="#5f00c7", bd=0, highlightthickness=0)
        self.reset_btn = Ctk.CTkButton(self, text="reset Image", width=100, command=self.reset_image, font=self.my_font)
        self.create_folder_btn = Ctk.CTkButton(self,image=self.create_folder_image, text="", width=100, command=self.create_new_folder, font=self.my_font)
        self.open_folder_btn = Ctk.CTkButton(self, image=self.open_folder_image,text="", width=100, command=self.open_new_folder, font=self.my_font)
        self.save_img_btn = Ctk.CTkButton(self, text="Save Image", width=100, command=self.save_and_load, font=self.my_font)
        self.delete_img_btn = Ctk.CTkButton(self, text="Delete Image", width=100, command=self.delete_img, font=self.my_font)
        self.skip_time_btn = Ctk.CTkButton(self, text="Jump", width=100, command=self.skip_time, font=self.my_font)
        self.new_source_btn = Ctk.CTkButton(self,image=self.source_folder_image, text="", width=100, command=self.source_folder_dialog, font=self.my_font)
        self.preview_canvas = Ctk.CTkCanvas(self,background="#5f00c7", bd=0, highlightthickness=0)
        self.start_mode = Ctk.StringVar(value=system_code.data_mode)
        self.choose_mode = Ctk.CTkSegmentedButton(self, values= system_code.data_modes,
                                                            variable=self.start_mode,
                                                            command=self.set_data_mode,width=64,font=self.my_font)

        # big_canvas
        self.big_canvas.bind("<ButtonPress-1>", self.on_press)
        self.big_canvas.bind("<B1-Motion>", self.on_drag)
        self.big_canvas.bind("<ButtonRelease-1>", self.on_release)
        self.after_idle(self.adjust_image_sizes_for_buttons)
        self.bind("<Configure>", self.on_resize)
        self.place_labeling()

    def set_data_mode(self, value):
        system_code.data_mode = value
        if system_code.data_mode == system_code.data_modes[0]:
            self.reset_canvas()
        else:
            if self.tk_cropped_image is not None:
                self.show_img(self.save_cropped)
            else:
                self.show_img(self.img_paths[self.index])

    def set_data_mode_binding(self, value = None):
        self.data_mode_index = 1 - self.data_mode_index
        system_code.data_mode = system_code.data_modes[self.data_mode_index]
        self.choose_mode.set(system_code.data_mode)

    def reset_image(self):
        self.reset_canvas()
        self.show_img(self.index)

    def adjust_image_sizes_for_buttons(self):
        """Adjust the image sizes based on the current height of the buttons and update them accordingly."""
        self.update_idletasks()  # Force Tkinter to finalize the window layout
        
        example_button_height = self.create_folder_btn.winfo_height()
        example_button_width = self.create_folder_btn.winfo_width()

        # Assuming buttons' height as a basis for square image size
        orientation_value = min(example_button_height, example_button_width) - 25
        
        # Ensure the orientation_value does not drop below a minimum size threshold
        orientation_value = max(orientation_value, 1)  # Prevents size from being 0 or negative

        # Your existing logic to resize and apply the images
        self.image_btn_size = (orientation_value, orientation_value)
        self.load_button_images()
        
        self.create_folder_btn.configure(image=self.create_folder_image)
        self.open_folder_btn.configure(image=self.open_folder_image)
        self.new_source_btn.configure(image=self.source_folder_image)

        self.create_folder_btn.image = self.create_folder_image
        self.open_folder_btn.image = self.open_folder_image
        self.new_source_btn.image = self.source_folder_image
        

    def load_button_images(self):
        """
        Load, resize button images to the specified size and convert them into a Tkinter-compatible format.
        """
        # Convert PIL images to  Ctk.CTkImage objects for compatibility with Tkinter/customtkinter
        self.create_folder_image =  Ctk.CTkImage(self.create_folder_raw_image , size =self.image_btn_size)
        self.open_folder_image =  Ctk.CTkImage(self.open_folder_raw_image, size =self.image_btn_size)
        self.source_folder_image =  Ctk.CTkImage(self.source_folder_raw_image, size =self.image_btn_size)
    
    def reset_canvas(self):
        """
        Adjust the reset_canvas method to remove all rectangles or drawing objects while preserving a specific image.
        """
        # Assuming the image you want to keep has a 'picture' tag
        # First, remove all items tagged with 'rectangle' from the canvas
        self.big_canvas.delete("rectangle")
        
        # Reset the list of rectangles since they've been removed from the canvas
        self.labeling_boxes.clear()
        self.labeling_box_index = 0
        
        # Reset state variables
        self.rect = None
        self.draw = None
        self.preview_canvas.delete("all")  # Assuming you also want to clear the preview canvas
        
    def save_and_load(self):       
        save_path = f"{self.active_output_path}/{self.save_index:03d}{system_code.img_format}"
        if system_code.data_mode == system_code.data_modes[1]:
            if len(self.labeling_boxes) == 0:
                self.labeled = False
            else:
                image = self.draw_rects(self.original_image)
            image.save(save_path)
            self.resolution = image.size
        elif self.tk_cropped_image is not None:
            self.save_cropped.save(save_path)
            self.resolution = self.save_cropped.size
        
        self.data_saver.append_to_json_file(save_path, self.resolution, self.labeled, self.labeling_boxes, self.mpenn_data[0])
        self.reset_canvas()
        os.remove(self.img_paths[self.index])
        self.index += 1   
        self.save_index += 1
        self.show_img(self.img_paths[self.index]) 
    
    def draw_rects(self, image):
        self.draw = ImageDraw.Draw(image)
        
        for rect_info in self.labeling_boxes:
            _, coordinates, color, thickness = rect_info
            
            # Draw the rectangle on the image
            if thickness <= 0:
                thickness = 1
            self.draw.rectangle(coordinates, outline=color, width=thickness)

        self.labeled = True
        return image

    def delete_img(self):
        os.remove(self.img_paths[self.index])
        self.index += 1   
        self.show_img(self.img_paths[self.index]) 

    def skip_time(self):
        for i in range(system_code.skipable_frames):
            os.remove(self.img_paths[self.index + i])
        self.index = self.index + system_code.skipable_frames
        self.show_img(self.img_paths[self.index]) 
        
    def show_img(self, to_displayed):
        self.preview_canvas.delete("all")
        self.tk_cropped_image = None
        # Check if 'to_displayed' is a PIL Image object
        if isinstance(to_displayed, Image.Image):
            self.original_image = to_displayed  # 'to_displayed' is already an image object, use it directly
        else:
            # It's assumed to be a file path, try to open as an image file
            try:
                self.original_image = Image.open(to_displayed)
            except Exception:
                return

        self.return_image_information()
        self.resize_image()
        self.display_image()
        
    def resize_image(self):
        """
        Resize the image to fit within the panel dimensions while maintaining aspect ratio.
        Note: This function should be called after the window is visible and has been updated, to ensure
        accurate dimensions are retrieved from 'big_canvas'.
        """

        self.big_canvas.update_idletasks()

        if self.original_image is None or self.big_canvas.winfo_width() <= 1 or self.big_canvas.winfo_height() <= 1:
            return 
        # Update idletasks to ensure the window's layout is processed and accurate dimensions are retrieved
        self.big_canvas.update_idletasks()

        panel_width, panel_height = self.big_canvas.winfo_width(), self.big_canvas.winfo_height()

        original_width, original_height = self.original_image.size
        aspect_ratio = original_width / original_height

        # Determine the best fit size for keeping the aspect ratio
        if panel_width / aspect_ratio <= panel_height:
            new_width, new_height = int(panel_width), int(panel_width / aspect_ratio)
        else:
            new_width, new_height = int(panel_height * aspect_ratio), int(panel_height)

        # Inside the resize_image method, after computing new_width and new_height
        if new_width <= 0 or new_height <= 0:
            return  # Consider logging this case or handling it appropriately

        self.displayed_image = self.original_image.resize((new_width, new_height), Image.LANCZOS)

        # Calculate the size_factor
        size_factor_width =  original_width / new_width
        size_factor_height =  original_height / new_height

        self.img_factor_x = size_factor_width
        self.img_factor_y = size_factor_height

    def display_image(self):
        """Display the image centered in the canvas."""
        if self.displayed_image is None:
            return

        if self.tk_image: 
            self.delete_current_rectangle("all")  # Clear the canvas before displaying a new image

        self.tk_image = ImageTk.PhotoImage(self.displayed_image)
        # Force an update to ensure current dimensions are fetched
        self.big_canvas.update_idletasks()
        canvas_width, canvas_height = self.big_canvas.winfo_width(), self.big_canvas.winfo_height()
        image_width, image_height = self.tk_image.width(), self.tk_image.height()

        # Calculate the center position
        x_offset = int((canvas_width - image_width) // 2)
        y_offset = int((canvas_height - image_height) // 2)

        self.big_canvas.create_image(x_offset, y_offset, anchor=Ctk.NW, image=self.tk_image)
        self.image_position = (x_offset, y_offset)  # Save the centered image position

    def on_press(self, event):
        self.start_x = event.x
        self.start_y = event.y
        # If there's an existing rectangle, remove it
        if self.rect and (system_code.data_mode == system_code.data_modes[0]):
            self.delete_current_rectangle(self.rect)
        # Create an initial 1x1 rectangle that will be adjusted in `on_drag`
        self.rect = self.big_canvas.create_rectangle(self.start_x, self.start_y, self.start_x+1, self.start_y+1, 
                                            outline=system_code.color if system_code.data_mode != "Resize" else "#5f00c7", 
                                            width=system_code.thickness // self.img_factor_x if system_code.data_mode != "Resize" else 1,
                                            tags=("rectangle",))

    def on_resize(self, event):
        """Handle window resize events with throttling."""
        if not self.resize_pending:
            self.resize_pending = True

        if self.after_id:
            self.after_cancel(self.after_id)
        self.adjust_image_sizes_for_buttons()
        self.after_id = self.after(100, self.perform_resize)
    
    def on_drag(self, event):
        if not self.dragging:
            # This marks the start of actual dragging, so we do the rectangle creation here
            self.dragging = True  # Set dragging flag to True to indicate dragging has started
            # If there's an existing rectangle from a previous operation, remove it
            if self.rect:
                self.delete_current_rectangle(self.rect)
            # Start a new rectangle
            self.rect = self.big_canvas.create_rectangle(self.start_x, self.start_y, self.start_x+1, self.start_y+1, 
                                                        outline=system_code.color if system_code.data_mode == system_code.data_modes[1] else "#5f00c7", 
                                                        width=system_code.thickness // self.img_factor_x if system_code.data_mode == system_code.data_modes[1] else 1,
                                                        tags=("rectangle",))
        
        self.end_x, self.end_y = event.x, event.y

        if system_code.data_mode == system_code.data_modes[0]:
            # Maintain 1:1 aspect ratio for resize mode
            delta_x = self.end_x - self.start_x
            delta_y = self.end_y - self.start_y
            delta = min(abs(delta_x), abs(delta_y))
            delta_x = delta if delta_x > 0 else -delta
            delta_y = delta if delta_y > 0 else -delta

            self.end_x, self.end_y = self.start_x + delta_x, self.start_y + delta_y
        # For other modes, the rectangle adjusts to current coordinates directly

        # Update rectangle size during drag
        self.big_canvas.coords(self.rect, self.start_x, self.start_y, self.end_x, self.end_y)

    def on_release(self, event):
        # Only proceed if dragging actually occurred
        if self.dragging:
            self.big_canvas.coords(self.rect, self.start_x, self.start_y, self.end_x, self.end_y)
            if self.start_x > self.end_x:
                temp = self.start_x
                self.start_x = self.end_x
                self.end_x = temp
            if self.start_y > self.end_y:
                temp = self.start_y
                self.start_y = self.end_y
                self.end_y = temp
            self.big_canvas.coords(self.rect, self.start_x, self.start_y, self.end_x, self.end_y)
            self.display_cropped_part()  # New line to display the cropped area
            self.dragging = False  # Reset dragging flag
            if system_code.data_mode == system_code.data_modes[1] and self.rect is not None:
                coordinates = [self.start_x, self.start_y, self.end_x, self.end_y]
                coordinates = self.subtract_image_from_canvas(coordinates)
                coordinates = self.multiply_array(coordinates, self.img_factor_x, self.img_factor_y)
                temp_array = [self.labeling_box_index,coordinates, system_code.color, system_code.thickness]
                self.labeling_boxes.append(temp_array)
                self.labeling_box_index += 1
        else:
            # Handle the case where a rectangle shouldn't have been started
            if self.rect:
                self.delete_current_rectangle(self.rect)
            self.rect = None

    def multiply_array(self, arr, num1, num2):
        """
        Multiplies each element of the array 'arr' by 'num'.

        Parameters:
        arr (list): The input list whose elements are to be multiplied.
        num (int): The number by which each element of the list is multiplied.

        Returns:
        list: A new list with each element of 'arr' multiplied by 'num'.
        """
        # Use list comprehension to multiply each element by num
        return arr[0] * num1, arr[1]* num2, arr[2]* num1, arr[3] * num2

    def subtract_image_from_canvas(self, arr):

        width = self.image_position[0]
        height = self.image_position[1]

        return arr[0] - width, arr[1] - height, arr[2] - width, arr[3] - height


    def display_cropped_part(self):
        if not self.validate_image_and_rectangle():
            return

        adj_coordinates = self.adjust_coordinates_for_border()
        if not self.validate_adjusted_coordinates(adj_coordinates):
            self.delete_current_rectangle(self.rect)
            return

        crop_coordinates = self.calculate_crop_coordinates(adj_coordinates)
        self.cropped_image = self.original_image.crop(crop_coordinates)
        self.display_resized_cropped_image(self.cropped_image)

    def adjust_coordinates_for_border(self):
        """Adjust rectangle's position by considering the image's position and border width."""
        adj_start_x = self.start_x - self.image_position[0] + (system_code.thickness / 2)
        adj_start_y = self.start_y - self.image_position[1] + (system_code.thickness / 2)
        adj_end_x = self.end_x - self.image_position[0] - (system_code.thickness / 2)
        adj_end_y = self.end_y - self.image_position[1] - (system_code.thickness / 2)
        return adj_start_x, adj_start_y, adj_end_x, adj_end_y

    def validate_adjusted_coordinates(self, coordinates):
        """Ensure logical and boundary adherence of adjusted coordinates."""
        adj_start_x, adj_start_y, adj_end_x, adj_end_y = coordinates
        if (adj_end_x <= adj_start_x) or (adj_end_y <= adj_start_y):
            # Negative width or height indicates a coordinate error
            return False
        
        if (adj_start_x < 0 or adj_start_y < 0 or 
                adj_end_x > self.displayed_image.width or 
                adj_end_y > self.displayed_image.height):
            # Coordinates exceed displayed image bounds
            return False
        
        return True

    def calculate_crop_coordinates(self, adj_coordinates):
        """Calculate crop coordinates on the original image, considering display scaling."""
        adj_start_x, adj_start_y, adj_end_x, adj_end_y = adj_coordinates
        
        # Scaling ratios, considering the displayed image might be scaled down/up to fit into the canvas
        x_scale = self.original_image.size[0] / self.displayed_image.width 
        y_scale = self.original_image.size[1] / self.displayed_image.height 

        # Apply scaling to adjust coordinates to match the original image size
        crop_start_x = max(0, int(adj_start_x * x_scale))
        crop_start_y = max(0, int(adj_start_y * y_scale))
        crop_end_x = min(self.original_image.size[0], int(adj_end_x * x_scale))
        crop_end_y = min(self.original_image.size[1], int(adj_end_y * y_scale))
        
        return crop_start_x, crop_start_y, crop_end_x, crop_end_y

    def display_resized_cropped_image(self, cropped_image):
        """Resize and display the cropped image on the preview canvas."""
        cropped_image_aspect_ratio = (cropped_image.width / cropped_image.height)

        preview_canvas_width = self.preview_canvas.winfo_width()
        preview_canvas_height = self.preview_canvas.winfo_height()
        if cropped_image_aspect_ratio > (preview_canvas_width / preview_canvas_height):
            new_width = preview_canvas_width
            new_height = int(preview_canvas_width / cropped_image_aspect_ratio)
        else:
            new_height = preview_canvas_height
            new_width = int(preview_canvas_height * cropped_image_aspect_ratio)

        resized_cropped_image = cropped_image.resize((new_width, new_height), Image.LANCZOS)
        self.save_cropped = resized_cropped_image
        self.tk_cropped_image = ImageTk.PhotoImage(resized_cropped_image)
        self.preview_canvas.delete("all")
        centered_x = (preview_canvas_width - new_width) / 2
        centered_y = (preview_canvas_height - new_height) / 2
        self.preview_canvas.create_image(centered_x, centered_y, anchor='nw', image=self.tk_cropped_image)
        self.preview_canvas.image = self.tk_cropped_image

    def perform_resize(self):
        """Perform the resize operation and clear cropped image and rectangle."""
        if self.resize_pending:
            self.resize_image()
            self.display_image()

            # Clear the rectangle on the main canvas if it exists
            self.reset_canvas()

            # Reset rectangle coordinates
            self.start_x, self.start_y, self.end_x, self.end_y = None, None, None, None
            
            # Optionally, if desired, clear the dragging states
            self.dragging = False


    def delete_current_rectangle(self, obj):
        """Remove the rectangle object if coordinates are invalid."""
        self.big_canvas.delete(obj)
        self.rect = None

    def validate_image_and_rectangle(self):
        if not self.displayed_image or not self.rect:
            return False
        return True

    def place_labeling(self):
        self.big_canvas.place(
                relx=0.025,
                rely=0.5,
                relwidth=0.775,
                relheight=0.95,
                anchor="w"
            ) 
        self.reset_btn.place(
                relx=0.975,
                rely=0.1,
                relwidth=0.15,
                relheight=0.06,
                anchor="e",
            ) 
        self.create_folder_btn.place(
                relx=0.975,
                rely=0.2,
                relwidth=0.06,
                relheight=0.06,
                anchor="e",
            ) 
        self.open_folder_btn.place(
                relx=0.825,
                rely=0.2,
                relwidth=0.06,
                relheight=0.06,
                anchor="w",
            ) 
        self.save_img_btn.place(
                relx=0.975,
                rely=0.3,
                relwidth=0.15,
                relheight=0.06,
                anchor="e",
            ) 
        self.delete_img_btn.place(
                relx=0.975,
                rely=0.4,
                relwidth=0.15,
                relheight=0.06,
                anchor="e",
            ) 
        self.skip_time_btn.place(
                relx=0.975,
                rely=0.5,
                relwidth=0.06,
                relheight=0.06,
                anchor="e",
            ) 
        self.new_source_btn.place(
                relx=0.825,
                rely=0.5,
                relwidth=0.06,
                relheight=0.06,
                anchor="w",
            ) 
        self.preview_canvas.place(
                relx=0.975,
                rely=0.7,
                relwidth=0.15,
                relheight=0.25,
                anchor="e",
                )  
        self.choose_mode.place(
                relx=0.975,
                rely=0.9,
                relwidth=0.15,
                relheight=0.06,
                anchor="e",
            ) 

    def hide_labeling(self):
        self.big_canvas.place_forget()
        self.reset_btn.place_forget()
        self.create_folder_btn.place_forget()
        self.open_folder_btn.place_forget()
        self.save_img_btn.place_forget()
        self.delete_img_btn.place_forget()
        self.skip_time_btn.place_forget()
        self.new_source_btn.place_forget()
        self.preview_canvas.place_forget()
        self.choose_mode.place_forget()