GENRATE-AI-DATA.pyw

import os
import sqlite3
import numpy as np
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import io
import re
from imageio.v2 import imread

# Constants
current_db_path = "DEEPDATA.db"
decode_ai_folder = "AI"
os.makedirs(decode_ai_folder, exist_ok=True)

# UI Colors
BG_COLOR = "#2b2d42"
BTN_COLOR = "#8d99ae"
BTN_TEXT_COLOR = "#edf2f4"
LISTBOX_BG = "#3d405b"
LISTBOX_FG = "#edf2f4"
HEADER_LEN = 4 * 8  # uint32 bit length for header

# Database Operations
def connect_db():
    """Create and connect to the SQLite database."""
    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS images (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT UNIQUE,
        image BLOB
    )
    """)
    conn.commit()
    conn.close()

# Listbox and Display
def alphanumeric_sort_key(item):
    """Sort key that splits strings into numeric and non-numeric parts."""
    return [int(text) if text.isdigit() else text.lower() for text in re.split('(\d+)', item)]

def load_images_from_sources():
    """Load images from both DecodeAI folder and DEEPDATA.db into the listbox, sorted alphanumerically."""
    image_list.delete(0, tk.END)

    images = []

    # Load images from DecodeAI folder
    for file_name in os.listdir(decode_ai_folder):
        if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
            images.append(file_name)

    # Load images from DEEPDATA.db
    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM images")
    db_images = cursor.fetchall()
    conn.close()

    # Add database images to the list
    for (file_name,) in db_images:
        if file_name not in images:  # Avoid duplicates
            images.append(file_name)

    # Sort images alphanumerically
    images.sort(key=alphanumeric_sort_key)

    # Populate the Listbox
    for file_name in images:
        image_list.insert(tk.END, file_name)

def view_selected_image():
    """Display the selected original image in the GUI, prioritizing DEEPDATA.db over DecodeAI folder."""
    selected_image = image_list.get(tk.ACTIVE)
    if not selected_image:
        messagebox.showwarning("No Selection", "No image selected.")
        return

    # Try to load the image from the database
    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT image FROM images WHERE name = ?", (selected_image,))
    row = cursor.fetchone()
    conn.close()

    if row:
        # If the image is found in the database
        try:
            image_data = row[0]
            image = Image.open(io.BytesIO(image_data))
            image.thumbnail((400, 400))
            original_tk_image = ImageTk.PhotoImage(image)
            original_label.config(image=original_tk_image)
            original_label.image = original_tk_image
            return
        except Exception as e:
            messagebox.showerror("Error", f"Could not display image from the database: {e}")

    # If not in database, try to load from DecodeAI folder
    original_path = os.path.join(decode_ai_folder, selected_image)
    if os.path.exists(original_path):
        try:
            original_image = Image.open(original_path)
            original_image.thumbnail((400, 400))
            original_tk_image = ImageTk.PhotoImage(original_image)
            original_label.config(image=original_tk_image)
            original_label.image = original_tk_image
        except Exception as e:
            messagebox.showerror("Error", f"Could not display the image from folder: {e}")
    else:
        messagebox.showinfo("Image Not Found", "Image not found in either the database or the folder.")

def save_selected_image():
    """Save the selected image from the database to the DecodeAI folder in PNG format."""
    selected_image = image_list.get(tk.ACTIVE)
    if not selected_image:
        messagebox.showwarning("No Selection", "No image selected.")
        return

    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT image FROM images WHERE name = ?", (selected_image,))
    result = cursor.fetchone()
    conn.close()

    if result:
        # Ensure the image name ends with .png for saving
        save_file_name = f"{os.path.splitext(selected_image)[0]}.png"
        save_path = os.path.join(decode_ai_folder, save_file_name)
        try:
            # Convert the BLOB data back to an image and save it
            image_data = io.BytesIO(result[0])
            image = Image.open(image_data)
            image.save(save_path, format="PNG")
            messagebox.showinfo("Saved", f"Image saved successfully to {save_path}")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to save the image: {e}")
    else:
        messagebox.showwarning("Not Found", "Could not find the image data.")

# New Functionality: Select Image from DB and Save to Folder
def add_image():
    """Add a new image to the database."""
    filepath = filedialog.askopenfilename(
        title="Select Image",
        filetypes=[("Image Files", "*.png;*.jpg;*.jpeg;*.bmp;*.gif")]
    )
    if not filepath:
        return

    name = os.path.basename(filepath)
    with open(filepath, 'rb') as file:
        image_data = file.read()

    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    try:
        cursor.execute("INSERT INTO images (name, image) VALUES (?, ?)", (name, image_data))
        conn.commit()
        #load_images()
        messagebox.showinfo("Success", f"Image '{name}' added successfully.")
    except sqlite3.IntegrityError:
        messagebox.showwarning("Duplicate", f"Image '{name}' already exists.")
    finally:
        conn.close()

def delete_image():
    """Delete the selected image."""
    selected_image = image_list.get(tk.ACTIVE)
    if not selected_image:
        messagebox.showwarning("No Selection", "No image selected.")
        return

    conn = sqlite3.connect(current_db_path)
    cursor = conn.cursor()
    cursor.execute("DELETE FROM images WHERE name = ?", (selected_image,))
    conn.commit()
    conn.close()
    #load_images()
    messagebox.showinfo("Deleted", f"Image '{selected_image}' deleted successfully.")

# Steganography Functions (Same as before)
def read_image(img_path):
    img = np.array(imread(img_path), dtype=np.uint8)
    return img.flatten(), img.shape

def decode_data(encoded_data):
    return np.bitwise_and(encoded_data, np.ones_like(encoded_data))

def write_file(file_path, file_bit_array):
    bytes_data = np.packbits(file_bit_array)
    with open(file_path, 'wb') as f:
        f.write(bytes_data)

def decode_hidden_image():
    """Decode and display the hidden image."""
    selected_image = image_list.get(tk.ACTIVE)
    if not selected_image:
        messagebox.showwarning("No Selection", "No image selected.")
        return

    encoded_image_path = os.path.join(decode_ai_folder, selected_image)
    encoded_data, shape_orig = read_image(encoded_image_path)
    data = decode_data(encoded_data)
    el_array = np.packbits(data[:HEADER_LEN])
    extracted_len = el_array.view(np.uint32)[0]
    hidden_data = data[HEADER_LEN:HEADER_LEN + extracted_len]

    # Save the decoded image
    hidden_image_path = os.path.join(decode_ai_folder, f"decoded_{selected_image}")
    write_file(hidden_image_path, hidden_data)

    # Preview the decoded image
    try:
        decoded_image = Image.open(hidden_image_path)
        decoded_image.thumbnail((400, 400))
        decoded_tk_image = ImageTk.PhotoImage(decoded_image)
        decoded_label.config(image=decoded_tk_image)
        decoded_label.image = decoded_tk_image
        messagebox.showinfo("Success", f"Hidden image saved and displayed from: {hidden_image_path}")
    except Exception as e:
        messagebox.showerror("Error", f"Failed to preview the decoded image: {e}")

# Refresh List Function
def refresh_list():
    """Refresh the image list in the Listbox."""
    load_images_from_sources()

# Main GUI
root = tk.Tk()
root.title("Najeeb AI Images Generator")
root.geometry("1000x700")
root.configure(bg=BG_COLOR)

connect_db()

# Listbox Frame
listbox_frame = tk.Frame(root, bg=BG_COLOR)
listbox_frame.pack(side=tk.LEFT, padx=10, pady=10)

scrollbar = tk.Scrollbar(listbox_frame, orient=tk.VERTICAL)
scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

image_list = tk.Listbox(listbox_frame, selectmode=tk.SINGLE, width=15, height=45, bg=LISTBOX_BG, fg=LISTBOX_FG, yscrollcommand=scrollbar.set)
image_list.pack(side=tk.LEFT)
scrollbar.config(command=image_list.yview)

# Button Frame (All buttons in one row)
button_frame = tk.Frame(root, bg=BG_COLOR)
button_frame.pack(pady=10)

view_button = tk.Button(button_frame, text="View Image", command=view_selected_image, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
view_button.pack(side=tk.LEFT, padx=5)

save_button = tk.Button(button_frame, text="Save Image", command=save_selected_image, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
save_button.pack(side=tk.LEFT, padx=5)

refresh_button = tk.Button(button_frame, text="Refresh List", command=refresh_list, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
refresh_button.pack(side=tk.LEFT, padx=5)

decode_button = tk.Button(button_frame, text="Decode Image", command=decode_hidden_image, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
decode_button.pack(side=tk.LEFT, padx=5)

# Add Button
add_button = tk.Button(button_frame, text="Add Image", command=add_image, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
add_button.pack(side=tk.LEFT, padx=5)

# Delete Button
delete_button= tk.Button(button_frame, text="Delete Image", command=delete_image, bg=BTN_COLOR, fg=BTN_TEXT_COLOR)
delete_button.pack(side=tk.LEFT, padx=5)

# Display Labels for Original and Decoded images side by side
image_frame = tk.Frame(root, bg=BG_COLOR)
image_frame.pack(pady=10)

original_label = tk.Label(image_frame)
original_label.pack(side=tk.LEFT, padx=10)

decoded_label = tk.Label(image_frame)
decoded_label.pack(side=tk.LEFT, padx=10)

# Load images into Listbox
load_images_from_sources()

root.mainloop()