dataImageGenerator.py

import os
import cv2
import numpy as np
import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
from PIL import Image, ImageEnhance, ImageTk
from random import randint, uniform
from zipfile import ZipFile
import threading
import sys
from datetime import datetime

class ImageAugmentationTool:
    def __init__(self, root):
        self.root = root
        self.root.title("Dataset Image Augmentation Generator")
        self.root.geometry("1200x800")
        self.root.resizable(True, True)
        
        self.input_folder = tk.StringVar()
        self.output_folder = tk.StringVar()
        self.progress_var = tk.DoubleVar()
        self.status_text = tk.StringVar(value="Ready to start...")
        self.total_images = 0
        self.processed_images = 0
        
        self.aug_options = {}
        
        self.setup_styles()
        
        self.create_widgets()
        
    def setup_styles(self):
        style = ttk.Style()
        style.theme_use('clam')
        style.configure('Title.TLabel', font=('Arial', 16, 'bold'))
        style.configure('Header.TLabel', font=('Arial', 12, 'bold'))
        style.configure('Success.TLabel', foreground='green')
        style.configure('Error.TLabel', foreground='red')
        style.configure('Accent.TButton', font=('Arial', 10, 'bold'))
        style.configure('TCheckbutton', indicatorsize=20)
        
    def create_widgets(self):
        main_canvas = tk.Canvas(self.root)
        scrollbar = ttk.Scrollbar(self.root, orient="vertical", command=main_canvas.yview)
        scrollable_frame = ttk.Frame(main_canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: main_canvas.configure(scrollregion=main_canvas.bbox("all"))
        )
        
        canvas_window = main_canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")

        def resize_canvas(event):
            main_canvas.itemconfig(canvas_window, width=event.width)
        main_canvas.bind("<Configure>", resize_canvas)
        
        main_canvas.configure(yscrollcommand=scrollbar.set)
        
        main_canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        def on_mousewheel(event):
            main_canvas.yview_scroll(int(-1*(event.delta/120)), "units")
        
        main_canvas.bind_all("<MouseWheel>", on_mousewheel)
        
        main_frame = ttk.Frame(scrollable_frame, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        title_label = ttk.Label(main_frame, text="Image Augmentation Tool", 
                               style='Title.TLabel')
        title_label.pack(pady=(0, 10))
        
        input_frame = ttk.LabelFrame(main_frame, text="Input Configuration", padding="5")
        input_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Label(input_frame, text="Select Class Folder:", 
                 style='Header.TLabel').grid(row=0, column=0, sticky=tk.W, pady=5)
        
        input_path_frame = ttk.Frame(input_frame)
        input_path_frame.grid(row=1, column=0, columnspan=3, sticky=tk.EW, pady=5)
        
        self.input_entry = ttk.Entry(input_path_frame, textvariable=self.input_folder, 
                                     width=60, font=('Arial', 10))
        self.input_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))
        
        ttk.Button(input_path_frame, text="Browse...", 
                  command=self.browse_input_folder).pack(side=tk.RIGHT)
        
        ttk.Label(input_frame, text="Supported formats: JPG, JPEG, PNG", 
                 foreground='gray').grid(row=2, column=0, sticky=tk.W, pady=(5, 0))
        
        output_frame = ttk.LabelFrame(main_frame, text="Output Configuration", padding="5")
        output_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Label(output_frame, text="Output Folder:", 
                 style='Header.TLabel').grid(row=0, column=0, sticky=tk.W, pady=5)
        
        output_path_frame = ttk.Frame(output_frame)
        output_path_frame.grid(row=1, column=0, columnspan=3, sticky=tk.EW, pady=5)
        
        self.output_entry = ttk.Entry(output_path_frame, textvariable=self.output_folder, 
                                      width=60, font=('Arial', 10))
        self.output_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 10))
        
        ttk.Button(output_path_frame, text="Browse...", 
                  command=self.browse_output_folder).pack(side=tk.RIGHT)
        
        self.use_default = tk.BooleanVar(value=True)
        
        options_frame = ttk.LabelFrame(main_frame, text="Augmentation Options (Select at least one)", 
                                      padding="5")
        options_frame.pack(fill=tk.X, pady=(0, 5))
        
        options = [
            ("1. Rotation (-3° to +3°)", "rotate"),
            ("2. Horizontal Flip", "flip"),
            ("3. Gaussian Noise", "noise"),
            ("4. Color Jitter", "color_jitter"),
            ("5. Random Crop (90%)", "crop"),
            ("6. Gaussian Blur", "blur"),
            ("7. Brightness Adjustment", "brightness"),
            ("8. Contrast Enhancement", "contrast"),
            ("9. Sharpening", "sharpen"),
            ("10. Elastic Deformation", "elastic"),
            ("11. Vertical Flip", "vflip"),
            ("12. Random Shear", "shear"),
            ("13. Random Zoom", "zoom"),
            ("14. Random Translation", "shift"),
            ("15. Perspective Transform", "perspective"),
            ("16. Cutout (Erasing)", "cutout"),
            ("17. Mixup (Blending)", "mixup"),
            ("18. CutMix", "cutmix"),
            ("19. CLAHE", "clahe"),
            ("20. Gamma Correction", "gamma"),
            ("21. Channel Shuffle", "channel_shuffle"),
            ("22. Random Grayscale", "grayscale"),
            ("23. Edge Enhancement", "edge_enhance"),
            ("24. Histogram Equalization", "hist_equal"),
            ("25. Motion Blur", "motion_blur"),
            ("26. Salt & Pepper Noise", "salt_pepper"),
            ("27. Random Occlusion", "occlusion"),
            ("28. Background Replace", "bg_replace"),
            ("29. Fourier Noise", "fourier_noise"),
            ("30. Posterize", "posterize")
        ]
        
        canvas_frame = ttk.Frame(options_frame)
        canvas_frame.pack(fill=tk.X, pady=5)
        
        inner_frame = ttk.Frame(canvas_frame)
        inner_frame.pack(fill=tk.X, expand=True)
        
        for i, (label, key) in enumerate(options):
            row = i // 5
            col = i % 5
            
            self.aug_options[key] = tk.BooleanVar(value=True)
            
            cb = tk.Checkbutton(inner_frame, 
                    text=label, 
                    variable=self.aug_options[key],
                    onvalue=True,
                    offvalue=False)
            cb.grid(row=row, column=col, sticky=tk.W, padx=5, pady=8)
        
        select_frame = ttk.Frame(options_frame)
        select_frame.pack(pady=5)
        
        ttk.Button(select_frame, text="Select All", 
                  command=self.select_all).pack(side=tk.LEFT, padx=5)
        ttk.Button(select_frame, text="Deselect All", 
                  command=self.deselect_all).pack(side=tk.LEFT, padx=5)
        
        self.selected_count_label = ttk.Label(options_frame, text="✓ Selected: 30 options", 
                                            font=('Arial', 10, 'bold'), foreground='green')
        self.selected_count_label.pack()
        
        for key in self.aug_options:
            self.aug_options[key].trace('w', self.update_selected_count)
        
        control_frame = ttk.Frame(main_frame)
        control_frame.pack(fill=tk.X, pady=(0, 15))
        
        self.process_btn = ttk.Button(control_frame, text="🚀 START AUGMENTATION", 
                                      command=self.start_augmentation, 
                                      style='Accent.TButton')
        self.process_btn.pack(side=tk.LEFT, padx=5, ipadx=10, ipady=5)
        
        ttk.Button(control_frame, text="Clear Output", 
                  command=self.clear_output).pack(side=tk.LEFT, padx=5)
        
        progress_frame = ttk.Frame(main_frame)
        progress_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.progress_bar = ttk.Progressbar(progress_frame, variable=self.progress_var, 
                                           maximum=100, length=400)
        self.progress_bar.pack(fill=tk.X, expand=True)
        
        self.status_label = ttk.Label(main_frame, textvariable=self.status_text, 
                                     font=('Arial', 10))
        self.status_label.pack(pady=(0, 10))
        
        download_frame = ttk.Frame(main_frame)
        download_frame.pack(fill=tk.X, pady=10)
        
        self.download_btn = ttk.Button(download_frame, text="📥 DOWNLOAD AUGMENTED DATASET (ZIP)", 
                                       command=self.download_zip, state=tk.DISABLED,
                                       style='Accent.TButton')
        self.download_btn.pack(pady=5, ipadx=10, ipady=5)
        
        log_frame = ttk.LabelFrame(main_frame, text="Process Log", padding="10")
        log_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        
        self.log_text = scrolledtext.ScrolledText(log_frame, height=8, width=80, 
                                                  font=('Courier', 9))
        self.log_text.pack(fill=tk.BOTH, expand=True)
        
    def toggle_output_entry(self):
        if self.use_default.get():
            self.output_entry.config(state='disabled')
            self.output_folder.set("")
        else:
            self.output_entry.config(state='normal')
    
    def browse_input_folder(self):
        folder = filedialog.askdirectory(title="Select Class Folder")
        if folder:
            self.input_folder.set(folder)
            self.log_message(f"Input folder selected: {folder}")
    
    def browse_output_folder(self):
        folder = filedialog.askdirectory(title="Select Output Folder")
        if folder:
            self.output_folder.set(folder)
            self.log_message(f"Output folder selected: {folder}")
    
    def log_message(self, message):
        timestamp = datetime.now().strftime("%H:%M:%S")
        self.log_text.insert(tk.END, f"[{timestamp}] {message}\n")
        self.log_text.see(tk.END)
        self.root.update()
    
    def clear_output(self):
        self.log_text.delete(1.0, tk.END)
        self.status_text.set("Ready to start...")
        self.progress_var.set(0)
        self.download_btn.config(state=tk.DISABLED)
    
    def select_all(self):
        for key in self.aug_options:
            self.aug_options[key].set(True)
    
    def deselect_all(self):
        for key in self.aug_options:
            self.aug_options[key].set(False)
    
    def update_selected_count(self, *args):
        count = sum(1 for var in self.aug_options.values() if var.get())
        self.selected_count_label.config(text=f"✓ Selected: {count} options")
    
    def get_selected_augmentations(self):
        return [key for key, var in self.aug_options.items() if var.get()]
    
    def rotate_image(self, img):
        angle = randint(-3, 3)
        h, w = img.shape[:2]
        M = cv2.getRotationMatrix2D((w//2, h//2), angle, 1)
        return cv2.warpAffine(img, M, (w, h))

    def flip_image(self, img):
        return cv2.flip(img, 1)

    def add_noise(self, img):
        noise = np.random.normal(0, 25, img.shape).astype(np.uint8)
        return cv2.add(img, noise)

    def color_jitter(self, img):
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        enhancer = ImageEnhance.Brightness(pil_img)
        pil_img = enhancer.enhance(uniform(0.7, 1.3))
        enhancer = ImageEnhance.Contrast(pil_img)
        pil_img = enhancer.enhance(uniform(0.7, 1.3))
        enhancer = ImageEnhance.Color(pil_img)
        pil_img = enhancer.enhance(uniform(0.7, 1.3))
        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

    def random_crop(self, img):
        h, w = img.shape[:2]
        crop_size = int(min(h, w) * 0.9)
        x = randint(0, w - crop_size)
        y = randint(0, h - crop_size)
        cropped = img[y:y+crop_size, x:x+crop_size]
        return cv2.resize(cropped, (w, h))

    def blur_image(self, img):
        return cv2.GaussianBlur(img, (5,5), 0)
    
    def adjust_brightness(self, img):
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        hsv = np.array(hsv, dtype=np.float64)
        hsv[:,:,2] = hsv[:,:,2] * uniform(0.5, 1.5)
        hsv[:,:,2][hsv[:,:,2] > 255] = 255
        hsv = np.array(hsv, dtype=np.uint8)
        return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    
    def adjust_contrast(self, img):
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        enhancer = ImageEnhance.Contrast(pil_img)
        pil_img = enhancer.enhance(uniform(0.5, 1.5))
        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    
    def sharpen_image(self, img):
        kernel = np.array([[-1,-1,-1],
                          [-1, 9,-1],
                          [-1,-1,-1]])
        return cv2.filter2D(img, -1, kernel)
    
    def elastic_deformation(self, img):
        h, w = img.shape[:2]
        dx = np.random.rand(h, w) * 10 - 5
        dy = np.random.rand(h, w) * 10 - 5
        x, y = np.meshgrid(np.arange(w), np.arange(h))
        x_map = (x + dx).astype(np.float32)
        y_map = (y + dy).astype(np.float32)
        return cv2.remap(img, x_map, y_map, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
    
    def vflip_image(self, img):
        return cv2.flip(img, 0)
    
    def shear_image(self, img):
        h, w = img.shape[:2]
        shear_factor = uniform(-0.2, 0.2)
        M = np.float32([[1, shear_factor, 0], [0, 1, 0]])
        return cv2.warpAffine(img, M, (w, h))
    
    def zoom_image(self, img):
        h, w = img.shape[:2]
        zoom_factor = uniform(0.8, 1.2)
        new_h, new_w = int(h * zoom_factor), int(w * zoom_factor)
        resized = cv2.resize(img, (new_w, new_h))
        if zoom_factor > 1:
            start_x = (new_w - w) // 2
            start_y = (new_h - h) // 2
            return resized[start_y:start_y+h, start_x:start_x+w]
        else:
            result = np.zeros_like(img)
            start_x = (w - new_w) // 2
            start_y = (h - new_h) // 2
            result[start_y:start_y+new_h, start_x:start_x+new_w] = resized
            return result
    
    def shift_image(self, img):
        h, w = img.shape[:2]
        tx = randint(-20, 20)
        ty = randint(-20, 20)
        M = np.float32([[1, 0, tx], [0, 1, ty]])
        return cv2.warpAffine(img, M, (w, h))
    
    def perspective_transform(self, img):
        h, w = img.shape[:2]
        pts1 = np.float32([[0,0], [w-1,0], [0,h-1], [w-1,h-1]])
        pts2 = np.float32([
            [w*0.05, h*0.05],
            [w*0.95, h*0.1],
            [w*0.1, h*0.95],
            [w*0.9, h*0.9]
        ])
        M = cv2.getPerspectiveTransform(pts1, pts2)
        return cv2.warpPerspective(img, M, (w, h))
    
    def cutout_aug(self, img):
        h, w = img.shape[:2]
        mask_size = min(h, w) // 4
        x = randint(0, w - mask_size)
        y = randint(0, h - mask_size)
        img_copy = img.copy()
        img_copy[y:y+mask_size, x:x+mask_size] = 0
        return img_copy
    
    def mixup_aug(self, img):
        return img
    
    def cutmix_aug(self, img):
        return img
    
    def clahe_aug(self, img):
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge([l, a, b])
        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
    
    def gamma_correction(self, img):
        gamma = uniform(0.7, 1.3)
        inv_gamma = 1.0 / gamma
        table = np.array([(i / 255.0) ** inv_gamma * 255 for i in range(256)]).astype(np.uint8)
        return cv2.LUT(img, table)
    
    def channel_shuffle(self, img):
        channels = cv2.split(img)
        perm = np.random.permutation(3)
        return cv2.merge([channels[perm[0]], channels[perm[1]], channels[perm[2]]])
    
    def grayscale_aug(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    
    def edge_enhance(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        sobel = np.sqrt(sobelx**2 + sobely**2)
        sobel = np.uint8(np.clip(sobel, 0, 255))
        return cv2.cvtColor(sobel, cv2.COLOR_GRAY2BGR)
    
    def hist_equalize(self, img):
        img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
        img_yuv[:,:,0] = cv2.equalizeHist(img_yuv[:,:,0])
        return cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
    
    def motion_blur(self, img):
        size = 15
        kernel_motion_blur = np.zeros((size, size))
        kernel_motion_blur[int((size-1)/2), :] = np.ones(size)
        kernel_motion_blur = kernel_motion_blur / size
        return cv2.filter2D(img, -1, kernel_motion_blur)
    
    def salt_pepper_noise(self, img):
        row, col, ch = img.shape
        s_vs_p = 0.5
        amount = 0.02
        out = img.copy()
        num_salt = np.ceil(amount * img.size * s_vs_p)
        coords = [np.random.randint(0, i-1, int(num_salt)) for i in [row, col]]
        out[coords[0], coords[1], :] = 255
        num_pepper = np.ceil(amount * img.size * (1. - s_vs_p))
        coords = [np.random.randint(0, i-1, int(num_pepper)) for i in [row, col]]
        out[coords[0], coords[1], :] = 0
        return out
    
    def occlusion_aug(self, img):
        h, w = img.shape[:2]
        patch_size = min(h, w) // 3
        x = randint(0, w - patch_size)
        y = randint(0, h - patch_size)
        img_copy = img.copy()
        random_patch = np.random.randint(0, 255, (patch_size, patch_size, 3), dtype=np.uint8)
        img_copy[y:y+patch_size, x:x+patch_size] = random_patch
        return img_copy
    
    def bg_replace_aug(self, img):
        return img
    
    def fourier_noise(self, img):
        f = np.fft.fft2(img)
        fshift = np.fft.fftshift(f)
        magnitude_spectrum = 20*np.log(np.abs(fshift))
        noise = np.random.normal(0, 10, f.shape)
        f_noisy = f + noise
        f_ishift = np.fft.ifftshift(f_noisy)
        img_back = np.fft.ifft2(f_ishift)
        img_back = np.abs(img_back).astype(np.uint8)
        return img_back
    
    def posterize_aug(self, img):
        bits = randint(3, 6)
        shift = 8 - bits
        return (img >> shift) << shift
    
    def apply_augmentations(self, img):
        augmented_images = []
        selected = self.get_selected_augmentations()
        
        aug_functions = {
            'rotate': self.rotate_image,
            'flip': self.flip_image,
            'noise': self.add_noise,
            'color_jitter': self.color_jitter,
            'crop': self.random_crop,
            'blur': self.blur_image,
            'brightness': self.adjust_brightness,
            'contrast': self.adjust_contrast,
            'sharpen': self.sharpen_image,
            'elastic': self.elastic_deformation,
            'vflip': self.vflip_image,
            'shear': self.shear_image,
            'zoom': self.zoom_image,
            'shift': self.shift_image,
            'perspective': self.perspective_transform,
            'cutout': self.cutout_aug,
            'mixup': self.mixup_aug,
            'cutmix': self.cutmix_aug,
            'clahe': self.clahe_aug,
            'gamma': self.gamma_correction,
            'channel_shuffle': self.channel_shuffle,
            'grayscale': self.grayscale_aug,
            'edge_enhance': self.edge_enhance,
            'hist_equal': self.hist_equalize,
            'motion_blur': self.motion_blur,
            'salt_pepper': self.salt_pepper_noise,
            'occlusion': self.occlusion_aug,
            'bg_replace': self.bg_replace_aug,
            'fourier_noise': self.fourier_noise,
            'posterize': self.posterize_aug
        }
        
        for aug_name in selected:
            if aug_name in aug_functions:
                try:
                    aug_img = aug_functions[aug_name](img.copy())
                    augmented_images.append(aug_img)
                except Exception as e:
                    self.log_message(f"Warning: {aug_name} augmentation failed: {str(e)}")
                    augmented_images.append(img.copy())
        
        return augmented_images
    
    def start_augmentation(self):
        if not self.input_folder.get():
            messagebox.showwarning("Warning", "Please select an input folder first!")
            return
        
        if not self.get_selected_augmentations():
            messagebox.showwarning("Warning", "Please select at least one augmentation option!")
            return
        
        thread = threading.Thread(target=self.process_images)
        thread.daemon = True
        thread.start()
    
    def process_images(self):
        self.process_btn.config(state=tk.DISABLED)
        self.download_btn.config(state=tk.DISABLED)
        
        try:
            class_folder = self.input_folder.get()
            class_name = os.path.basename(class_folder)
            
            if self.use_default.get() or not self.output_folder.get():
                output_folder = os.path.join("augmented_dataset", class_name)
            else:
                output_folder = self.output_folder.get()
            
            os.makedirs(output_folder, exist_ok=True)
            
            image_files = [f for f in os.listdir(class_folder) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            
            if not image_files:
                self.log_message("No images found in the folder!")
                messagebox.showwarning("Warning", "No images found in the selected folder!")
                return
            
            selected_augs = self.get_selected_augmentations()
            total_files = len(image_files)
            processed = 0
            augmented_count = 0
            
            self.log_message(f"Found {total_files} images to process")
            self.log_message(f"Selected augmentations: {len(selected_augs)} options")
            self.log_message(f"Output folder: {output_folder}")
            
            for img_file in image_files:
                img_path = os.path.join(class_folder, img_file)
                img = cv2.imread(img_path)
                
                if img is None:
                    self.log_message(f"Failed to load: {img_file}")
                    continue
                
                augmented_imgs = self.apply_augmentations(img)
                
                base_name = os.path.splitext(img_file)[0]
                for i, aug_img in enumerate(augmented_imgs):
                    save_path = os.path.join(output_folder, f"{base_name}_aug{i+1}.jpg")
                    cv2.imwrite(save_path, aug_img)
                    augmented_count += 1
                
                processed += 1
                progress = (processed / total_files) * 100
                self.progress_var.set(progress)
                self.status_text.set(f"Processing: {processed}/{total_files} images")
                self.log_message(f"Processed: {img_file} -> {len(augmented_imgs)} augmented images")
                self.root.update()
            
            self.status_text.set(f"✅ Completed! Generated {augmented_count} augmented images")
            self.log_message(f"Total augmented images saved: {augmented_count}")
            self.log_message(f"Output saved to: {output_folder}")
            
            self.last_output_folder = output_folder
            self.download_btn.config(state=tk.NORMAL)
            
            messagebox.showinfo("Success", f"Augmentation completed!\n\n"
                              f"Original images: {total_files}\n"
                              f"Selected augmentations: {len(selected_augs)}\n"
                              f"Total generated: {augmented_count} images\n\n"
                              f"You can now download the ZIP file.")
            
        except Exception as e:
            self.log_message(f"ERROR: {str(e)}")
            messagebox.showerror("Error", f"An error occurred:\n{str(e)}")
        
        finally:
            self.process_btn.config(state=tk.NORMAL)
    
    def download_zip(self):
        if not hasattr(self, 'last_output_folder'):
            messagebox.showwarning("Warning", "No augmented data available to zip!")
            return
        
        try:
            initial_file = f"augmented_{os.path.basename(self.last_output_folder)}.zip"
            zip_path = filedialog.asksaveasfilename(
                defaultextension=".zip",
                filetypes=[("ZIP files", "*.zip"), ("All files", "*.*")],
                initialfile=initial_file,
                title="Save ZIP file as"
            )
            
            if not zip_path:
                return
            
            self.status_text.set("Creating zip file...")
            self.root.update()
            
            with ZipFile(zip_path, 'w') as zipf:
                for root, dirs, files in os.walk(self.last_output_folder):
                    for file in files:
                        file_path = os.path.join(root, file)
                        arcname = os.path.relpath(file_path, self.last_output_folder)
                        zipf.write(file_path, arcname=arcname)
            
            self.log_message(f"ZIP file created: {zip_path}")
            self.status_text.set("✅ ZIP file created successfully!")
            
            if messagebox.askyesno("Success", f"ZIP file saved to:\n{zip_path}\n\nOpen containing folder?"):
                folder_path = os.path.dirname(zip_path)
                if os.name == 'nt':  
                    os.startfile(folder_path)
                elif os.name == 'posix':  
                    import subprocess
                    subprocess.run(['open', folder_path] if sys.platform == 'darwin' else ['xdg-open', folder_path])
            
        except Exception as e:
            self.log_message(f"ERROR creating ZIP: {str(e)}")
            messagebox.showerror("Error", f"Failed to create ZIP:\n{str(e)}")

def main():
    root = tk.Tk()
    app = ImageAugmentationTool(root)
    
    root.update_idletasks()
    width = min(1200, root.winfo_screenwidth() - 100)
    height = min(800, root.winfo_screenheight() - 100)
    x = (root.winfo_screenwidth() // 2) - (width // 2)
    y = (root.winfo_screenheight() // 2) - (height // 2)
    root.geometry(f'{width}x{height}+{x}+{y}')
    
    root.mainloop()

if __name__ == "__main__":
    main()