app.py

import os
import zipfile
import shutil
import time
from PIL import Image, ImageDraw
from io import BytesIO
import io
from rembg import remove
import gradio as gr
from concurrent.futures import ThreadPoolExecutor
from transformers import AutoModelForImageSegmentation, pipeline
import numpy as np
import pandas as pd
import json
import requests
from dotenv import load_dotenv
import torch
from torchvision import transforms
from functools import lru_cache
import cv2
import pillow_avif
import threading
from collections import Counter
from transformers.configuration_utils import PretrainedConfig
if not hasattr(PretrainedConfig, "get_text_config"):
    PretrainedConfig.get_text_config = lambda self: None

stop_event = threading.Event()

# Load environment variables
load_dotenv()
PHOTOROOM_API_KEY = os.getenv("PHOTOROOM_API_KEY", "e98517e5e68a1a2eee49b130c2bcef05c1faec42")

_birefnet_model = None
_birefnet_transform = None
_birefnet_hr_model = None
_birefnet_hr_transform = None

@lru_cache(maxsize=1)
def get_birefnet_model():
    global _birefnet_model, _birefnet_transform
    if _birefnet_model is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        _birefnet_model = AutoModelForImageSegmentation.from_pretrained(
            'ZhengPeng7/BiRefNet',
            trust_remote_code=True,
            torch_dtype=torch.float32
        ).to(device)
        if not hasattr(_birefnet_model.config, "get_text_config"):
            _birefnet_model.config.get_text_config = lambda: None
        _birefnet_model.eval()
        _birefnet_transform = transforms.Compose([
            transforms.Resize((1024, 1024)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    return _birefnet_model, _birefnet_transform

def get_birefnet_hr_model():
    global _birefnet_hr_model, _birefnet_hr_transform
    if _birefnet_hr_model is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        _birefnet_hr_model = AutoModelForImageSegmentation.from_pretrained(
            'ZhengPeng7/BiRefNet_HR',
            trust_remote_code=True,
            torch_dtype=torch.float32
        ).to(device)
        if not hasattr(_birefnet_hr_model.config, "get_text_config"):
            _birefnet_hr_model.config.get_text_config = lambda: None
        _birefnet_hr_model.eval()
        _birefnet_hr_transform = transforms.Compose([
            transforms.Resize((2048, 2048)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    return _birefnet_hr_model, _birefnet_hr_transform

def remove_background_rembg(input_path):
    print(f"Removing background using rembg for image: {input_path}")
    with open(input_path, 'rb') as f:
        input_image = f.read()
    out_data = remove(input_image)
    return Image.open(io.BytesIO(out_data)).convert("RGBA")

def remove_background_bria(input_path):
    print(f"Removing background using bria for image: {input_path}")
    device = 0 if torch.cuda.is_available() else -1
    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device=device)
    result = pipe(input_path)
    if isinstance(result, list) and len(result) > 0 and "mask" in result[0]:
        mask = result[0]["mask"]
    else:
        mask = result
    if mask.mode != "RGBA":
        mask = mask.convert("RGBA")
    return mask

def remove_background_birefnet(input_path):
    try:
        model, transform_image = get_birefnet_model()
        device = next(model.parameters()).device
        image = Image.open(input_path).convert("RGB")
        input_tensor = transform_image(image).unsqueeze(0).to(device)
        with torch.no_grad():
            try:
                preds = model(input_tensor)[-1].sigmoid()
                pred_mask = preds[0].squeeze().cpu()
            except RuntimeError as e:
                if 'out of memory' in str(e):
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    input_tensor = input_tensor.cpu()
                    model = model.cpu()
                    preds = model(input_tensor)[-1].sigmoid()
                    pred_mask = preds[0].squeeze()
                    model = model.to(device)
                else:
                    raise e
        mask_pil = transforms.ToPILImage()(pred_mask)
        mask_resized = mask_pil.resize(image.size, Image.LANCZOS)
        result = image.copy()
        result.putalpha(mask_resized)
        result_array = np.array(result)
        alpha = result_array[:, :, 3]
        _, alpha = cv2.threshold(alpha, 248, 255, cv2.THRESH_BINARY)
        kernel_small = np.ones((3, 3), np.uint8)
        kernel_medium = np.ones((5, 5), np.uint8)
        kernel_large = np.ones((9, 9), np.uint8)
        alpha = cv2.GaussianBlur(alpha, (5, 5), 0)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_OPEN, kernel_small, iterations=3)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, kernel_medium, iterations=3)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, kernel_large, iterations=2)
        alpha = cv2.bilateralFilter(alpha, 9, 100, 100)
        alpha = cv2.medianBlur(alpha, 5)
        _, alpha = cv2.threshold(alpha, 248, 255, cv2.THRESH_BINARY)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_OPEN, kernel_small, iterations=2)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, kernel_small, iterations=2)
        edges = cv2.Canny(alpha, 100, 200)
        alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, kernel_medium, iterations=1)
        alpha = cv2.subtract(alpha, edges)
        result_array[:, :, 3] = alpha
        result = Image.fromarray(result_array)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return result
    except Exception as e:
        print(f"Error in remove_background_birefnet: {str(e)}")
        import traceback
        traceback.print_exc()
        raise

def remove_background_birefnet_2(input_path):
    model, transform_image = get_birefnet_model()
    device = next(model.parameters()).device
    image = Image.open(input_path).convert("RGB")
    input_tensor = transform_image(image).unsqueeze(0).to(device)
    with torch.no_grad():
        try:
            preds = model(input_tensor)[-1].sigmoid()
            pred_mask = preds[0].squeeze().cpu()
        except RuntimeError as e:
            if 'out of memory' in str(e):
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                input_tensor = input_tensor.cpu()
                model = model.cpu()
                preds = model(input_tensor)[-1].sigmoid()
                pred_mask = preds[0].squeeze()
                model = model.to(device)
            else:
                raise e
    mask_pil = transforms.ToPILImage()(pred_mask)
    mask_resized = mask_pil.resize(image.size, Image.LANCZOS)
    result = image.copy()
    result.putalpha(mask_resized)
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return result

def remove_background_birefnet_hr(input_path):
    try:
        model, transform_img = get_birefnet_hr_model()
        device = next(model.parameters()).device
        img = Image.open(input_path).convert("RGB")
        t_in = transform_img(img).unsqueeze(0).to(device)
        with torch.no_grad():
            preds = model(t_in)[-1].sigmoid()
            mask = preds[0].squeeze().cpu()
        mask_pil = transforms.ToPILImage()(mask).resize(img.size, Image.LANCZOS)
        out = img.copy()
        out.putalpha(mask_pil)
        return out.convert("RGBA")
    except Exception as e:
        print(f"remove_background_birefnet_hr: {e}")
        return None

def remove_background_photoroom(input_path):
    if input_path.lower().endswith('.avif'):
        input_path = convert_avif(input_path, input_path.rsplit('.', 1)[0] + '.png', 'PNG')    
    if not PHOTOROOM_API_KEY:
        raise ValueError("Photoroom API key missing.")
    url = "https://sdk.photoroom.com/v1/segment"
    headers = {"Accept": "image/png, application/json", "x-api-key": PHOTOROOM_API_KEY}
    with open(input_path, "rb") as f:
        resp = requests.post(url, headers=headers, files={"image_file": f})
    if resp.status_code != 200:
        raise Exception(f"PhotoRoom API error: {resp.status_code} - {resp.text}")
    return Image.open(BytesIO(resp.content)).convert("RGBA")

def remove_background_none(input_path):
    print(f"Removing background using none for image: {input_path}")
    return Image.open(input_path).convert("RGBA")

def get_dominant_color(image):
    tmp = image.convert("RGBA")
    tmp.thumbnail((100, 100))
    ccount = Counter(tmp.getdata())
    return ccount.most_common(1)[0][0]

def convert_avif(input_path, output_path, output_format='PNG'):
    with Image.open(input_path) as img:
        if output_format == 'JPG':
            img.convert("RGB").save(output_path, "JPEG")  # Convert to JPG (RGB mode)
        else:
            img.save(output_path, "PNG")  # Convert to PNG

    return output_path

def rotate_image(image, rotation, direction):
    if not image or rotation == "None":
        return image
    if rotation == "90 Degrees":
        angle = 90 if direction == "Clockwise" else -90
    elif rotation == "180 Degrees":
        angle = 180
    else:
        angle = 0
    return image.rotate(angle, expand=True)

def flip_image(image):
    return image.transpose(Image.FLIP_LEFT_RIGHT)

def get_bounding_box_with_threshold(image, threshold=10):
    arr = np.array(image)
    alpha = arr[:, :, 3]
    rows = np.any(alpha > threshold, axis=1)
    cols = np.any(alpha > threshold, axis=0)
    r_idx = np.where(rows)[0]
    c_idx = np.where(cols)[0]
    if r_idx.size == 0 or c_idx.size == 0:
        return None
    top, bottom = r_idx[0], r_idx[-1]
    left, right = c_idx[0], c_idx[-1]
    if left < right and top < bottom:
        return (left, top, right, bottom)
    else:
        return None

## === NEW ==
def position_logic_old(image_path, canvas_size, padding_top, padding_right, padding_bottom, padding_left, 
                   use_threshold=True, bg_method=None, is_person=False, 
                   snap_to_top=False, snap_to_bottom=False, snap_to_left=False, snap_to_right=False):
    """
    Position and resize an image on a canvas based on snapping, cropped sides, and birefnet logic.
    
    Args:
        image_path (str): Path to the input image.
        canvas_size (tuple): Target canvas size (width, height).
        padding_top, padding_right, padding_bottom, padding_left (int): Padding on each side.
        use_threshold (bool): Use threshold-based bounding box detection.
        bg_method (str): Background removal method ('birefnet', 'birefnet_2', etc.).
        is_person (bool): Treat as a person image (snaps to bottom by default).
        snap_to_top, snap_to_bottom, snap_to_left, snap_to_right (bool): Snap to respective sides.

    Returns:
        tuple: (log, resized_image, x_position, y_position)
    """
    # Load and prepare the image
    image = Image.open(image_path).convert("RGBA")
    log = []
    x, y = 0, 0
    
    # Get bounding box and crop
    if use_threshold:
        bbox = get_bounding_box_with_threshold(image, threshold=10)  # Assume this function exists
    else:
        bbox = image.getbbox()
    
    if bbox:
        # Detect cropped sides
        width, height = image.size
        cropped_sides = []
        tolerance = 30
        if any(image.getpixel((x, 0))[3] > tolerance for x in range(width)):
            cropped_sides.append("top")
        if any(image.getpixel((x, height-1))[3] > tolerance for x in range(width)):
            cropped_sides.append("bottom")
        if any(image.getpixel((0, y))[3] > tolerance for y in range(height)):
            cropped_sides.append("left")
        if any(image.getpixel((width-1, y))[3] > tolerance for y in range(height)):
            cropped_sides.append("right")
        if cropped_sides:
            log.append({"info": f"The following sides may contain cropped objects: {', '.join(cropped_sides)}"})
        else:
            log.append({"info": "The image is not cropped."})
        
        image = image.crop(bbox)
        log.append({"action": "crop", "bbox": [str(bbox[0]), str(bbox[1]), str(bbox[2]), str(bbox[3])]})
        
        # Setup variables
        target_width, target_height = canvas_size
        aspect_ratio = image.width / image.height
        
        # Determine active snaps
        snaps_active = []
        if padding_top == 0 or snap_to_top:
            snaps_active.append("top")
        if padding_bottom == 0 or snap_to_bottom or is_person:
            snaps_active.append("bottom")
        if padding_left == 0 or snap_to_left:
            snaps_active.append("left")
        if padding_right == 0 or snap_to_right:
            snaps_active.append("right")
        
        # Snap handling
        if snaps_active:
            if "top" in snaps_active and "bottom" in snaps_active:
                # Dual vertical snap: fill height
                new_height = target_height
                new_width = int(new_height * aspect_ratio)
                image = image.resize((new_width, new_height), Image.LANCZOS)
                y = 0
                if "left" in snaps_active:
                    x = 0
                elif "right" in snaps_active:
                    x = target_width - new_width
                else:
                    x = (target_width - new_width) // 2
                log.append({"action": "resize_snap_vertical", "new_width": str(new_width), "new_height": str(new_height)})
                log.append({"action": "position_snap_vertical", "x": str(x), "y": str(y)})
            elif "left" in snaps_active and "right" in snaps_active:
                # Dual horizontal snap: fill width
                new_width = target_width
                new_height = int(new_width / aspect_ratio)
                image = image.resize((new_width, new_height), Image.LANCZOS)
                x = 0
                if "top" in snaps_active:
                    y = 0
                elif "bottom" in snaps_active:
                    y = target_height - new_height
                else:
                    y = (target_height - new_height) // 2
                log.append({"action": "resize_snap_horizontal", "new_width": str(new_width), "new_height": str(new_height)})
                log.append({"action": "position_snap_horizontal", "x": str(x), "y": str(y)})
            else:
                # Original snap logic
                available_width = target_width
                available_height = target_height
                if "left" not in snaps_active:
                    available_width -= padding_left
                if "right" not in snaps_active:
                    available_width -= padding_right
                if "top" not in snaps_active:
                    available_height -= padding_top
                if "bottom" not in snaps_active:
                    available_height -= padding_bottom
                
                if aspect_ratio < 1:  # Portrait
                    new_height = available_height
                    new_width = int(new_height * aspect_ratio)
                    if new_width > available_width:
                        new_width = available_width
                        new_height = int(new_width / aspect_ratio)
                else:  # Landscape
                    new_width = available_width
                    new_height = int(new_width / aspect_ratio)
                    if new_height > available_height:
                        new_height = available_height
                        new_width = int(new_height * aspect_ratio)
                
                image = image.resize((new_width, new_height), Image.LANCZOS)
                if "left" in snaps_active:
                    x = 0
                elif "right" in snaps_active:
                    x = target_width - new_width
                else:
                    x = padding_left + (available_width - new_width) // 2
                if "top" in snaps_active:
                    y = 0
                elif "bottom" in snaps_active:
                    y = target_height - new_height
                else:
                    y = padding_top + (available_height - new_height) // 2
                log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
                log.append({"action": "position", "x": str(x), "y": str(y)})
        else:
            # No snaps: use cropped sides logic
            if len(cropped_sides) == 4:
                # All sides cropped: center crop to fit
                if aspect_ratio > 1:
                    new_height = target_height
                    new_width = int(new_height * aspect_ratio)
                    left = (new_width - target_width) // 2
                    image = image.resize((new_width, new_height), Image.LANCZOS)
                    image = image.crop((left, 0, left + target_width, target_height))
                else:
                    new_width = target_width
                    new_height = int(new_width / aspect_ratio)
                    top = (new_height - target_height) // 2
                    image = image.resize((new_width, new_height), Image.LANCZOS)
                    image = image.crop((0, top, target_width, top + target_height))
                x, y = 0, 0
                log.append({"action": "center_crop_resize", "new_size": f"{target_width}x{target_height}"})
            elif not cropped_sides:
                # No cropping: fit within padding
                new_height = target_height - padding_top - padding_bottom
                new_width = int(new_height * aspect_ratio)
                if new_width > target_width - padding_left - padding_right:
                    new_width = target_width - padding_left - padding_right
                    new_height = int(new_width / aspect_ratio)
                image = image.resize((new_width, new_height), Image.LANCZOS)
                x = (target_width - new_width) // 2
                y = target_height - new_height - padding_bottom
                log.append({"action": "resize", "new_width": str(new_width), "new_height": str(new_height)})
                log.append({"action": "position", "x": str(x), "y": str(y)})
            else:
                # Partial cropping: implement specific cases as needed
                # For simplicity, assume centering as a fallback
                new_width = target_width - padding_left - padding_right
                new_height = int(new_width / aspect_ratio)
                if new_height > target_height - padding_top - padding_bottom:
                    new_height = target_height - padding_top - padding_bottom
                    new_width = int(new_height * aspect_ratio)
                image = image.resize((new_width, new_height), Image.LANCZOS)
                x = (target_width - new_width) // 2
                y = (target_height - new_height) // 2
                log.append({"action": "resize_partial_crop", "new_width": str(new_width), "new_height": str(new_height)})
                log.append({"action": "position_partial_crop", "x": str(x), "y": str(y)})
        
        # Birefnet override
        if bg_method in ['birefnet', 'birefnet_2']:
            target_width = min(canvas_size[0] // 2, image.width)
            target_height = min(canvas_size[1] // 2, image.height)
            if aspect_ratio > 1:
                new_width = target_width
                new_height = int(new_width / aspect_ratio)
            else:
                new_height = target_height
                new_width = int(new_height * aspect_ratio)
            image = image.resize((new_width, new_height), Image.LANCZOS)
            x = (canvas_size[0] - new_width) // 2
            y = (canvas_size[1] - new_height) // 2
            log.append({"action": "birefnet_resize", "new_size": f"{new_width}x{new_height}", "position": f"{x},{y}"})
    
    return log, image, x, y

def position_logic_none(image, canvas_size):
    target_width, target_height = canvas_size
    aspect_ratio = image.width / image.height
    
    # Berikan margin di semua sisi (misalnya 50px dari setiap tepi)
    margin = 50
    available_width = target_width - (2 * margin)
    available_height = target_height - (2 * margin)
    
    # Scale factor untuk memperkecil gambar (85% dari ukuran available space)
    scale_factor = 0.85
    max_width = int(available_width * scale_factor)
    max_height = int(available_height * scale_factor)
    
    # Tentukan ukuran yang tepat dengan mempertahankan aspect ratio
    # dan memastikan gambar tidak terlalu besar (diperkecil dulu)
    if aspect_ratio > 1:  # landscape
        new_width = min(max_width, target_width - (2 * margin))
        new_height = int(new_width / aspect_ratio)
        if new_height > max_height:
            new_height = max_height
            new_width = int(new_height * aspect_ratio)
    else:  # portrait
        new_height = min(max_height, target_height - (2 * margin))
        new_width = int(new_height * aspect_ratio)
        if new_width > max_width:
            new_width = max_width
            new_height = int(new_width / aspect_ratio)
    
    # Resize gambar dengan ukuran baru (lebih kecil)
    image = image.resize((new_width, new_height), Image.LANCZOS)
    
    # Posisi tengah canvas 
    x = (target_width - new_width) // 2
    y = (target_height - new_height) // 2
    
    print(f"Image scaled down and centered: original_size={image.size}, new_size={new_width}x{new_height}, position=({x},{y}), margin={margin}px")
    log = [{"action": "scale_down_and_center", "new_size": f"{new_width}x{new_height}", "position": f"{x},{y}", "margin": f"{margin}px"}]
    return log, image, x, y

# ------------------ Qwen 2.5VL Inference Functions & Model Loading ------------------
import base64
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
import tempfile
import os
import base64

def encode_image(image_path):
    try:
        with open(image_path, "rb") as f:
            image_bytes = f.read()
        return base64.b64encode(image_bytes).decode('utf-8')
    except Exception as e:
        print(f"Error in encode_image: {str(e)}")
        raise

def classify_image(image_path, unique_items):
    try:
        image = Image.open(image_path).convert("RGB")
        image = image.resize((224, 224), Image.LANCZOS)
        
        print(f"Classifying image: {image_path} (resized to {image.size})")
        prompt = (
            f"Classify this image into one of these categories: {', '.join(unique_items)}. "
            f"Be sensitive to sizes of an object, e.g. 'small' or 'medium' or 'large', especially for bags. "
            f"If a hand is detected, only pick classifications that mention 'hand', however if it\'s a human, only pick classifications which mentioned 'human'. "
            f"Return only the classification word, nothing else."
        )
        
        # Save resized image to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
            image.save(temp_file.name, format='PNG')
            temp_image_path = temp_file.name
        
        # Get raw classification from API with retry logic
        classification_result = inference_with_api(temp_image_path, prompt)
        print(f"Raw API response for {image_path}: '{classification_result}'")
        
        # Clean up temporary file
        os.unlink(temp_image_path)
        
        # Parse and match the classification result
        classification_result = classification_result.strip().lower()
        for item in unique_items:
            if item.lower() in classification_result:
                print(f"Matched classification for {image_path}: '{item}'")
                return item
        
        print(f"No matching classification found in response: '{classification_result}'. Expected one of: {unique_items}")
        return None
    
    except Exception as e:
        print(f"Error during classification for {image_path}: {str(e)}")
        return None

def analyze_image_for_snap_settings(image_path):
    """
    Menganalisis gambar menggunakan Qwen untuk menentukan pengaturan snap yang tepat
    """
    try:
        prompt = (
            "Analyze this product/model/person image and determine if it should be flush against any edges of the canvas.\n\n"
            "For each edge (top, bottom, left, right), determine if the image should have padding=0 for that edge based on these specific rules:\n\n"
            "1. snap_bottom=true: If it's a person/model (almost always), or if the bottom of the product is cropped or should align with bottom edge\n\n"
            "2. snap_left=true: If the left side of a HAND or PRODUCT is cut off or flush against the edge, or if the hand or product is shown from side view facing left\n\n"
            "3. snap_right=true: If the right side a HAND or PRODUCT is cut off or flush against the edge, or if the hand or product is shown from side view facing right\n\n"
            "4. snap_top=true: If it's a person/model (almost always) or if the top of the product is cut off or should align with top edge\n\n"
            "Pay special attention to product orientation: side views often need snap_left or snap_right, while front/back views may not.\n\n"
            "EXAMPLES:\n"
            "- For a swimwear model standing and showing profile view: {\"snap_top\": false, \"snap_right\": false, \"snap_bottom\": true, \"snap_left\": true}\n"
            "- For a handbag shown from the side with handle at top: {\"snap_top\": false, \"snap_right\": false, \"snap_bottom\": true, \"snap_left\": true}\n"
            "- For a bikini bottom piece shown from front: {\"snap_top\": false, \"snap_right\": false, \"snap_bottom\": false, \"snap_left\": false}\n"
            "- For a swimsuit top on a model shown from side: {\"snap_top\": false, \"snap_right\": true, \"snap_bottom\": false, \"snap_left\": false}\n\n"
            "Common combinations:\n"
            "- For people/models, usually snap_bottom=true, snap_top=true and sometimes snap_left or snap_right depending on pose\n"
            "- For bags shown from side, use both snap_bottom=true and either snap_left=true or snap_right=true\n"
            "- For footwear shown from side, consider snap_bottom=true and either snap_left=true or snap_right=true\n"
            "- For items cropped on multiple sides, set all appropriate snap values to true\n\n"
            "Return ONLY a valid JSON in this exact format: {\"snap_top\": true/false, \"snap_right\": true/false, \"snap_bottom\": true/false, \"snap_left\": true/false}"
        )
        
        # Save image to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
            image = Image.open(image_path)
            image.save(temp_file.name, format='PNG')
            temp_image_path = temp_file.name
        
        # Get analysis from API
        analysis_result = inference_with_api(temp_image_path, prompt)
        print(f"Raw analysis response for {image_path}: '{analysis_result}'")
        
        # Clean up temporary file
        os.unlink(temp_image_path)
        
        # Parse JSON from the response
        try:
            # Coba parse langsung dulu
            try:
                snap_settings = json.loads(analysis_result)
                if all(key in snap_settings for key in ["snap_top", "snap_right", "snap_bottom", "snap_left"]):
                    print(f"Direct JSON parsing successful for {image_path}: {snap_settings}")
                    return snap_settings
            except:
                pass  # Lanjut ke regex jika direct parsing gagal
                
            # Mencari JSON dalam respons menggunakan regex
            import re
            json_match = re.search(r'(\{.*?\})', analysis_result, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
                snap_settings = json.loads(json_str)
                print(f"Parsed snap settings for {image_path}: {snap_settings}")
                return snap_settings
            else:
                print(f"No JSON found in response for {image_path}")
                return None
        except json.JSONDecodeError as e:
            print(f"Failed to parse JSON from response for {image_path}: {e}")
            return None
    
    except Exception as e:
        print(f"Error during snap setting analysis for {image_path}: {str(e)}")
        return None

def analyze_image_pattern(image_path):
    """
    Analyzes image patterns to determine snap settings based on cropped sides, whitespace, and content distribution.
    """
    try:
        # Initialize snap settings
        settings = {
            'snap_top': False,
            'snap_right': False,
            'snap_bottom': False,
            'snap_left': False
        }

        # Load and convert image to RGBA
        img = Image.open(image_path).convert("RGBA")
        img_np = np.array(img)
        height, width = img_np.shape[:2]
        aspect_ratio = height / width

        # Define mask for foreground pixels (alpha > 128)
        mask = img_np[:, :, 3] > 128

        # **Detect cropped sides** (foreground pixels within 5 pixels of edges)
        top_cropped = np.any(mask[:5, :])
        bottom_cropped = np.any(mask[-5:, :])
        left_cropped = np.any(mask[:, :5])
        right_cropped = np.any(mask[:, -5:])

        # **Detect big whitespace** (regions with >80% pixels having alpha < 128)
        top_whitespace = np.mean(img_np[:height//4, :, 3] < 128) > 0.8
        bottom_whitespace = np.mean(img_np[height - height//4:, :, 3] < 128) > 0.8
        left_whitespace = np.mean(img_np[:, :width//4, 3] < 128) > 0.8
        right_whitespace = np.mean(img_np[:, width - width//4:, 3] < 128) > 0.8

        # **Apply user-specified rules**
        if top_whitespace and bottom_whitespace and top_cropped and bottom_cropped:
            settings['snap_top'] = True
            settings['snap_bottom'] = True
        if top_whitespace and bottom_whitespace and left_whitespace and top_cropped and bottom_cropped and left_cropped:
            settings['snap_top'] = True
            settings['snap_bottom'] = True
            settings['snap_left'] = True
        if top_whitespace and bottom_whitespace and right_whitespace and top_cropped and bottom_cropped and right_cropped:
            settings['snap_top'] = True
            settings['snap_bottom'] = True
            settings['snap_right'] = True
        if bottom_whitespace and not top_whitespace and not left_whitespace and not right_whitespace and bottom_cropped and not top_cropped and not left_cropped and not right_cropped:
            settings['snap_bottom'] = True
        if top_whitespace and not bottom_whitespace and not left_whitespace and not right_whitespace and top_cropped and not bottom_cropped and not left_cropped and not right_cropped:
            settings['snap_top'] = True

        # **Additional logic from previous code**
        # Set snap_bottom for portrait images if not already set
        # Analyze vertical distribution for snap_top if not already set
        if not settings['snap_bottom']:
            bottom_foreground_ratio = np.mean(mask[height - height//4:, :])
            if bottom_foreground_ratio > 0.05:  # More than 5% foreground pixels in top quarter
                settings['snap_bottom'] = True

        # Analyze horizontal distribution if left or right snaps are not set
        if not (settings['snap_left'] or settings['snap_right']):
            horizontal_dist = np.sum(mask, axis=0)
            left_sum = np.sum(horizontal_dist[:width//3])
            right_sum = np.sum(horizontal_dist[2*width//3:])
            if left_sum > 1.5 * right_sum:
                settings['snap_left'] = True
            elif right_sum > 1.5 * left_sum:
                settings['snap_right'] = True

        # Analyze vertical distribution for snap_top if not already set
        if not settings['snap_top'] and aspect_ratio > 1.5:
                    settings['snap_top'] = True

        return settings

    except Exception as e:
        print(f"Error in analyze_image_pattern: {e}")
        return {
            'snap_top': False,
            'snap_right': False,
            'snap_bottom': False,
            'snap_left': False
        }
        
# ------------------ Modified process_single_image ------------------
def process_single_image(
    image_path,
    output_folder,
    bg_method,
    canvas_size_name,
    output_format,
    bg_choice,
    custom_color,
    watermark_path=None,
    twibbon_path=None,
    rotation=None,
    direction=None,
    flip=False,
    use_old_position=True,
    sheet_data=None,               # DataFrame with sheet data (if provided)
    use_qwen=False,
    snap_to_bottom=False,
    snap_to_top=False,
    snap_to_left=False,
    snap_to_right=False,
    auto_snap=False                # Tambahan parameter untuk mengaktifkan auto snap
):
    filename = os.path.basename(image_path)
    base_no_ext, ext = os.path.splitext(filename.lower())
    add_padding_line = False

    # ================== FULL SET OF CANVAS SIZE IFS ==================
    # Handle custom canvas size as tuple
    if isinstance(canvas_size_name, tuple):
        canvas_size = canvas_size_name
        padding_top = 100
        padding_right = 100
        padding_bottom = 100
        padding_left = 100
    elif canvas_size_name == 'Rox- Columbia & Keen':
        canvas_size = (1080, 1080)
        padding_top = 112
        padding_right = 126
        padding_bottom = 116
        padding_left = 126
    elif canvas_size_name == 'Jansport- Zalora':
        canvas_size = (762, 1100)
        padding_top = 108
        padding_right = 51
        padding_bottom = 202
        padding_left = 51
    elif canvas_size_name == 'Shopify & Lazada- Herschel':
        canvas_size = (1080, 1080)
        padding_top = 200
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Zalora- Herschel & Hedgren':
        canvas_size = (762, 1100)
        padding_top = 51
        padding_right = 51
        padding_bottom = 202
        padding_left = 51
    elif canvas_size_name == 'Jansport & Bratpack & Travelon & Hedgren- Lazada':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Jansport-Human- Lazada':
        canvas_size = (1080, 1080)
        padding_top = 72
        padding_right = 200
        padding_bottom = 180
        padding_left = 200        
    elif canvas_size_name == 'DC- Shopify':
        canvas_size = (1000, 1000)
        padding_top = 50
        padding_right = 80
        padding_bottom = 50
        padding_left = 80
    elif canvas_size_name == 'DC- S&L':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'ROX- Hydroflask-Shopify':
        canvas_size = (1080, 1080)
        padding_top = 112
        padding_right = 280
        padding_bottom = 116
        padding_left = 274
    elif canvas_size_name == 'Delsey- Lazada & Shopee':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 72
        padding_bottom = 180
        padding_left = 72
    elif canvas_size_name == 'Grind- Keen- Shopify':
        canvas_size = (1124, 1285)
        padding_top = 32
        padding_right = 127
        padding_bottom = 80
        padding_left = 132
    elif canvas_size_name == 'Bratpack- Gregory & DBTK- Shopify':
        canvas_size = (900, 1200)
        padding_top = 72
        padding_right = 66
        padding_bottom = 63
        padding_left = 66
    elif canvas_size_name == 'Columbia- Lazada':
        canvas_size = (1080, 1080)
        padding_top = 72
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Topo Design MP- Tiktok':
        canvas_size = (1080, 1080)
        padding_top = 200
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Columbia- Shopee & Zalora':
        canvas_size = (762, 1100)
        padding_top = 51
        padding_right = 51
        padding_bottom = 202
        padding_left = 51
    elif canvas_size_name == 'RTR- Columbia- Shopify':
        canvas_size = (1100, 737)
        padding_top = 38
        padding_right = 31
        padding_bottom = 39
        padding_left = 31
    elif canvas_size_name == 'columbia.psd':
        canvas_size = (730 , 610)
        padding_top = 29
        padding_right = 105
        padding_bottom = 36
        padding_left = 105
    elif canvas_size_name == 'jansport-dotcom':
        canvas_size = (1126, 1307)
        padding_top = 50
        padding_right = 50
        padding_bottom = 55
        padding_left = 50
    elif canvas_size_name == 'jansport-tiktok':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'quiksilver-lazada':
        canvas_size = (1080, 1080)
        padding_top = 200
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'quiksilver-shopee':
        canvas_size = (1080, 1080)
        padding_top = 200
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'grind':
        canvas_size = (1124, 1285)
        padding_top = 32
        padding_right = 127
        padding_bottom = 80
        padding_left = 132
    elif canvas_size_name == 'Allbirds- Shopee & Rockport':
        canvas_size = (1080, 1080)
        if base_no_ext.endswith(("_05")):
            padding_top = 440
        else:
            padding_top = 180
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Allbirds- Shopify':
        canvas_size = (1124, 1285)
        if base_no_ext.endswith("_05"):
            padding_top = 700
        else:
            padding_top = 175
        padding_right = 127
        padding_bottom = 80
        padding_left = 132
    elif canvas_size_name == 'Billabong- S&L':
        canvas_size = (1080, 1080)
        padding_top = 72
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'Quiksilver- Shopify':
        canvas_size = (1000, 1000)
        padding_top = 50
        padding_right = 80
        padding_bottom = 256
        padding_left = 80
    elif canvas_size_name == 'TTC-Shopify & Tiktok':
        canvas_size = (2800, 3201)
        padding_top = 392
        padding_right = 50
        padding_bottom = 50
        padding_left = 50
    elif canvas_size_name == 'Hydroflask- Shopee':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 315
        padding_bottom = 180
        padding_left = 315
    elif canvas_size_name == 'Hydroflask- Shopify':
        canvas_size = (1000, 1100)
        padding_top = 46
        padding_right = 348
        padding_bottom = 46
        padding_left = 348
    elif canvas_size_name == 'WT- New- Shopify':
        canvas_size = (2917, 3750)
        padding_top = 629
        padding_right = 608
        padding_bottom = 450
        padding_left = 600
    elif canvas_size_name == 'Roxy-Shopee':
        canvas_size = (1080, 1080)
        padding_top = 72
        padding_right = 200
        padding_bottom = 180
        padding_left = 200 
    elif canvas_size_name == 'Skechers':
        canvas_size = (3000, 3000)
        padding_top = 0
        padding_right = 0
        padding_bottom = 0
        padding_left = 0         
    elif canvas_size_name == 'Grind- Knockaround- Shopify':
        canvas_size = (1124, 1285)
        if base_no_ext.endswith("_03"):
            padding_top = 175
        else:
            padding_top = 694
        if base_no_ext.endswith("_03"):
            padding_bottom = 79
        else:
            padding_bottom = 204            
        padding_right = 127
        padding_left = 132   
    elif canvas_size_name == 'Sledgers-Lazada':
        canvas_size = (1080, 1080)
        padding_top = 420
        padding_right = 200
        padding_bottom = 180
        padding_left = 200 
    elif canvas_size_name == 'Aetrex-Lazada':
        canvas_size = (1080, 1080)
        padding_top = 180
        padding_right = 200
        padding_bottom = 180
        padding_left = 200
    elif canvas_size_name == 'primer-sale.psd':
        canvas_size = (700, 800)
        padding_top = 13
        padding_right = 13
        padding_bottom = 100
        padding_left = 12   
    elif canvas_size_name == 'TUMI-Shopify':
        canvas_size = (620, 750)
        padding_top = 297
        padding_right = 30
        padding_bottom = 56
        padding_left = 30            
    else:
        canvas_size = (1080, 1080)
        padding_top = 100
        padding_right = 100
        padding_bottom = 100
        padding_left = 100

# Classification and padding override
    classification_result = None
    
    # Logika Auto Snap yang independen dari klasifikasi
    if auto_snap:
        try:
            print(f"Auto snap enabled, analyzing image for optimal snap settings")
            
            # 1. Aplikasikan aturan preset terlebih dahulu (berdasarkan nama file)
            preset_settings = preset_snap_rules(filename, image_path)
            print(f"Preset snap settings for {filename}: {preset_settings}")
            
            # Jika tidak ada preset khusus yang cocok (semua False), lanjut ke metode lain
            if not any(preset_settings.values()):
                print(f"No preset rules match for {filename}, proceeding to pattern analysis")
                
                # 2. Analisis pola visual gambar (pendekatan berbasis computer vision)
                pattern_settings = analyze_image_pattern(image_path)
                print(f"Pattern analysis results for {filename}: {pattern_settings}")
                
                # Jika pattern analysis berhasil mendeteksi setidaknya satu snap
                if any(pattern_settings.values()):
                    # Gunakan hasil pattern analysis
                    snap_to_top = pattern_settings.get("snap_top", snap_to_top)
                    snap_to_right = pattern_settings.get("snap_right", snap_to_right)
                    snap_to_bottom = pattern_settings.get("snap_bottom", snap_to_bottom)
                    snap_to_left = pattern_settings.get("snap_left", snap_to_left)
                    print(f"Using pattern analysis results: top={snap_to_top}, right={snap_to_right}, bottom={snap_to_bottom}, left={snap_to_left}")
                else:
                    # 3. Jika pattern analysis tidak memberikan hasil, gunakan AI
                    print(f"Pattern analysis inconclusive for {filename}, attempting AI analysis")
                    snap_settings = analyze_image_for_snap_settings(image_path)
                    
                    if snap_settings:
                        # Validasi hasil snap settings
                        valid_snap = True
                        for key, value in snap_settings.items():
                            if not isinstance(value, bool):
                                print(f"Warning: Invalid value for {key}: {value}, expected boolean")
                                valid_snap = False
                        
                        # Hanya terapkan jika hasil valid
                        if valid_snap:
                            # Override manual snap settings dengan hasil analisis
                            snap_to_top = snap_settings.get("snap_top", snap_to_top)
                            snap_to_right = snap_settings.get("snap_right", snap_to_right)
                            snap_to_bottom = snap_settings.get("snap_bottom", snap_to_bottom)
                            snap_to_left = snap_settings.get("snap_left", snap_to_left)
                            print(f"AI snap settings applied: top={snap_to_top}, right={snap_to_right}, bottom={snap_to_bottom}, left={snap_to_left}")
                        else:
                            print(f"Invalid AI snap settings detected, using manual settings instead")
                    else:
                        print(f"Unable to determine optimal snap settings with AI, using manual settings instead")
            else:
                # Gunakan preset settings jika ada
                snap_to_top = preset_settings.get("snap_top", snap_to_top)
                snap_to_right = preset_settings.get("snap_right", snap_to_right)
                snap_to_bottom = preset_settings.get("snap_bottom", snap_to_bottom)
                snap_to_left = preset_settings.get("snap_left", snap_to_left)
                print(f"Using preset snap settings: top={snap_to_top}, right={snap_to_right}, bottom={snap_to_bottom}, left={snap_to_left}")
            
            # Final settings logging
            if snap_to_top:
                print(f"Auto snap: Setting top padding to 0 for {filename}")
            if snap_to_right:
                print(f"Auto snap: Setting right padding to 0 for {filename}")
            if snap_to_bottom:
                print(f"Auto snap: Setting bottom padding to 0 for {filename}")
            if snap_to_left:
                print(f"Auto snap: Setting left padding to 0 for {filename}")
                
        except Exception as e:
            print(f"Error during auto snap analysis for {filename}: {e}")
            print(f"Using manual snap settings due to auto snap error in {filename}.")
    
    # Klasifikasi untuk padding (tidak mempengaruhi auto snap)
    if use_qwen and sheet_data is not None:  # Only perform classification if toggle is on and sheet data exists
        try:
            unique_items = sheet_data['Classification'].str.strip().str.lower().unique().tolist()
            if not unique_items:
                print(f"No unique items found in sheet for {filename}. Using default padding.")
            else:
                print(f"Unique items for classification of {filename}: {unique_items}")
                classification_result = classify_image(image_path, unique_items)
                if classification_result is not None:
                    classification = classification_result.strip().lower()
                    print(f"Final classification for {filename}: '{classification}'")
                    if any(term in classification.lower() for term in ["human", "person", "model"]):
                        print(f"Person detected, setting bottom padding to 0 for {filename}")
                        snap_to_bottom = True
                    
                    matched_row = sheet_data[sheet_data['Classification'].str.strip().str.lower() == classification]
                    if not matched_row.empty:
                        row = matched_row.iloc[0]
                        padding_top = int(row['padding_top'])
                        padding_bottom = int(row['padding_bottom'])
                        padding_left = int(row['padding_left'])
                        padding_right = int(row['padding_right'])
                        print(f"Padding overridden for {filename}: top={padding_top}, bottom={padding_bottom}, left={padding_left}, right={padding_right}\n")
                    else:
                        print(f"No match found in sheet for classification '{classification}' in {filename}. Using default padding.\n")
                else:
                    print(f"Classification failed for {filename}. Using default padding.")
        except Exception as e:
            print(f"Error during classification for {filename}: {e}")
            print(f"Using default padding due to classification error in {filename}.")
    else:
        print(f"Qwen classification not used or no sheet data for {filename}. Using default padding.")

    padding_used = {
        "top": int(padding_top),
        "bottom": int(padding_bottom),
        "left": int(padding_left),
        "right": int(padding_right)
    }

    # Background removal and positioning (unchanged)
    if stop_event.is_set():
        print("Stop event triggered, no processing.")
        return None, None, None  # Return None for classification too

    print(f"Processing image: {filename}")
    original_img = Image.open(image_path).convert("RGBA")
    
    # Parse custom color to ensure it's in the correct format
    custom_color = parse_color(custom_color)
    if bg_method == 'rembg':
        mask = remove_background_rembg(image_path)
    elif bg_method == 'bria':
        mask = remove_background_bria(image_path)
    elif bg_method == 'photoroom':
        mask = remove_background_photoroom(image_path)
    elif bg_method == 'birefnet':
        mask = remove_background_birefnet(image_path)
        if not mask:
            return None, None
    elif bg_method == 'birefnet_2':
        mask = remove_background_birefnet_2(image_path)
        if not mask:
            return None, None
    elif bg_method == 'birefnet_hr':
        mask = remove_background_birefnet_hr(image_path)
        if not mask:
            return None, None
    elif bg_method == 'none':
        mask = original_img.copy()
        final_width, final_height = canvas_size
        orig_w, orig_h = mask.size
        threshold = 250
        rgb_mask = mask.convert('RGB')
        np_mask = np.array(rgb_mask)
        def is_column_white(col):
            return np.all(np_mask[:, col, 0] >= threshold) and np.all(np_mask[:, col, 1] >= threshold) and np.all(np_mask[:, col, 2] >= threshold)
        left_crop = 0
        while left_crop < orig_w and is_column_white(left_crop):
            left_crop += 1
        right_crop = orig_w - 1
        while right_crop > 0 and is_column_white(right_crop):
            right_crop -= 1
        if left_crop < right_crop:
            mask = mask.crop((left_crop, 0, right_crop + 1, orig_h))
    mask_array = np.array(mask)
    if bg_method == 'none':
        new_image_array = np.array(mask)
    else:
        new_image_array = np.array(original_img)
    new_image_array[:, :, 3] = mask_array[:, :, 3]
    image_with_no_bg = Image.fromarray(new_image_array)
    temp_image_path = os.path.join(output_folder, f"temp_{filename}")
    image_with_no_bg.save(temp_image_path, format='PNG')
    
    # Selalu gunakan position_logic_none untuk centering gambar
    # Kode snap masih disimpan untuk kompatibilitas
    if snap_to_left:
        print(f"Snap to Left active: Forcing padding_left = 0 (original: {padding_left})")
    if snap_to_right:
        print(f"Snap to Right active: Forcing padding_right = 0 (original: {padding_right})")
    if snap_to_top:
        print(f"Snap to Top active: Forcing padding_top = 0 (original: {padding_top})")
    if snap_to_bottom:
        print(f"Snap to Bottom active: Forcing padding_bottom = 0 (original: {padding_bottom})")
        
    # Gunakan position_logic_none untuk memastikan semua gambar diletakkan di tengah
    image = Image.open(temp_image_path)
    logs, cropped_img, x, y = position_logic_none(image, canvas_size)
    if bg_choice == 'white':
        canvas = Image.new("RGBA", canvas_size, "WHITE")
    elif bg_choice == 'custom':
        canvas = Image.new("RGBA", canvas_size, custom_color)
    elif bg_choice == 'dominant':
        dom_col = get_dominant_color(original_img)
        canvas = Image.new("RGBA", canvas_size, dom_col)
    else:
        canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0))
    canvas.paste(cropped_img, (x, y), cropped_img)
    logs.append({"action": "paste", "x": int(x), "y": int(y)})
    if flip:
        canvas = flip_image(canvas)
        logs.append({"action": "flip_horizontal"})
    if rotation != "None" and (rotation == "180 Degrees" or direction != "None"):
        if rotation == "90 Degrees":
            angle = 90 if direction == "Clockwise" else -90
        elif rotation == "180 Degrees":
            angle = 180
        else:
            angle = 0
        rotated_subject = cropped_img.rotate(angle, expand=True)
        if bg_choice == 'white':
            new_canvas = Image.new("RGBA", canvas_size, "WHITE")
        elif bg_choice == 'custom':
            new_canvas = Image.new("RGBA", canvas_size, custom_color)
        elif bg_choice == 'dominant':
            dom_col = get_dominant_color(original_img)
            new_canvas = Image.new("RGBA", canvas_size, dom_col)
        else:
            new_canvas = Image.new("RGBA", canvas_size, (0, 0, 0, 0))
            
        # Gunakan position_logic_none untuk rotated image juga
        _, rotated_sized_img, rotated_x, rotated_y = position_logic_none(rotated_subject, canvas_size)
        
        new_canvas.paste(rotated_sized_img, (rotated_x, rotated_y), rotated_sized_img)
        canvas = new_canvas
        logs.append({"action": "rotate_final_centered", "rotation": rotation, "direction": direction})
    out_ext = "jpg" if output_format == "JPG" else "png"
    out_filename = f"{os.path.splitext(filename)[0]}.{out_ext}"
    out_path = os.path.join(output_folder, out_filename)
    if (base_no_ext.endswith("_01") or base_no_ext.endswith("_1") or base_no_ext.endswith("_001")) and watermark_path:
        w_img = Image.open(watermark_path).convert("RGBA")
        canvas.paste(w_img, (0, 0), w_img)
        logs.append({"action": "add_watermark"})
    if twibbon_path:
        twb = Image.open(twibbon_path).convert("RGBA")
        canvas.paste(twb, (0, 0), twb)
        logs.append({"action": "twibbon"})
    if output_format == "JPG":
        canvas.convert("RGB").save(out_path, "JPEG")
    else:
        canvas.save(out_path, "PNG")
    os.remove(temp_image_path)
    print(f"Processed => {out_path}")
    return [(out_path, image_path)], logs, classification_result, padding_used

# ------------------ Modified process_images ------------------
def process_images(
    input_files,
    bg_method='rembg',
    watermark_path=None,
    twibbon_path=None,
    canvas_size='Rox- Columbia & Keen',
    output_format='PNG',
    bg_choice='transparent',
    custom_color="#ffffff",
    num_workers=4,
    rotation=None,
    direction=None,
    flip=False,
    use_old_position=True,
    progress=gr.Progress(),
    sheet_file=None,
    use_qwen=False,
    snap_to_bottom=False,
    snap_to_top=False,
    snap_to_left=False,
    snap_to_right=False,
    auto_snap=False
):
    stop_event.clear()
    start = time.time()
    if bg_method in ['birefnet', 'birefnet_2']:
        num_workers = 1
    out_folder = "processed_images"
    if os.path.exists(out_folder):
        shutil.rmtree(out_folder)
    os.makedirs(out_folder)
    procd = []
    origs = []
    all_logs = []
    classifications = {}

    # Load sheet file if provided
    sheet_data = None
    if sheet_file is not None:
        try:
            file_path = sheet_file.name if hasattr(sheet_file, "name") else sheet_file
            print(f"Attempting to load sheet file: {file_path}")
            if file_path.lower().endswith(".xlsx"):
                sheet_data = pd.read_excel(file_path)
            elif file_path.lower().endswith(".csv"):
                sheet_data = pd.read_csv(file_path)
            else:
                print(f"Unsupported file format for sheet: {file_path}")
            if sheet_data is not None:
                print(f"Sheet data loaded successfully with columns: {sheet_data.columns.tolist()}")
                # Validate required columns
                required_cols = {'Classification', 'padding_top', 'padding_bottom', 'padding_left', 'padding_right'}
                missing_cols = required_cols - set(sheet_data.columns)
                if missing_cols:
                    print(f"Warning: Missing required columns in sheet: {missing_cols}")
        except Exception as e:
            print(f"Error loading sheet file '{file_path}': {str(e)}")
            sheet_data = None

    # Input handling (unchanged)
    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
        tmp_in = "temp_input"
        if os.path.exists(tmp_in):
            shutil.rmtree(tmp_in)
        os.makedirs(tmp_in)
        with zipfile.ZipFile(input_files, 'r') as zf:
            zf.extractall(tmp_in)
        images = [os.path.join(tmp_in, f) for f in os.listdir(tmp_in) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp', '.tif', '.tiff', '.avif'))]
    elif isinstance(input_files, list):
        images = input_files
    else:
        images = [input_files]
    total = len(images)

    with ThreadPoolExecutor(max_workers=num_workers) as exe:
        future_map = {
            exe.submit(
                process_single_image,
                path,
                out_folder,
                bg_method,
                canvas_size,
                output_format,
                bg_choice,
                custom_color,
                watermark_path,
                twibbon_path,
                rotation,
                direction,
                flip,
                use_old_position,
                sheet_data,
                use_qwen,
                snap_to_bottom,
                snap_to_top,
                snap_to_left,
                snap_to_right,
                auto_snap
            ): path for path in images
        }
        for idx, fut in enumerate(future_map):
            if stop_event.is_set():
                print("Stop event triggered.")
                break
            try:
                result, log, classification, padding_used = fut.result()
                if result:
                    procd.extend(result)
                    origs.append(future_map[fut])
                    all_logs.append({os.path.basename(future_map[fut]): log})
                    classifications[os.path.basename(future_map[fut])] = {
                        "classification": classification if classification else "N/A",
                        "padding": padding_used
                    }
                progress((idx + 1) / total, f"{idx + 1}/{total} processed")
            except Exception as e:
                print(f"Error processing {future_map[fut]}: {str(e)}")

    # Save classifications (unchanged)
    with open(os.path.join(out_folder, "classifications.json"), "w") as cf:
        json.dump(classifications, cf, indent=2)
    zip_out = "processed_images.zip"
    with zipfile.ZipFile(zip_out, 'w') as zf:
        for outf, _ in procd:
            zf.write(outf, os.path.basename(outf))
    with open(os.path.join(out_folder, "process_log.json"), "w") as lf:
        json.dump(all_logs, lf, indent=2)
    elapsed = time.time() - start
    print(f"Done in {elapsed:.2f}s")
    return origs, procd, zip_out, elapsed, classifications

# ------------------ Gradio UI Setup ------------------
import gradio as gr
from concurrent.futures import ThreadPoolExecutor

def gradio_interface(
    input_files,
    bg_method,
    watermark,
    twibbon,
    canvas_size,
    output_format,
    bg_choice,
    custom_color,
    num_workers,
    rotation=None,
    direction=None,
    flip=False,
    sheet_file=None,
    use_qwen= False,  # sheet file input
    snap_to_bottom=False,
    snap_to_top=False,
    snap_to_left=False,
    snap_to_right=False,
    auto_snap=False
):
    if bg_method in ['birefnet', 'birefnet_2', 'birefnet_hr']:
        num_workers = min(num_workers, 2)
    progress = gr.Progress()
    watermark_path = watermark.name if watermark else None
    twibbon_path = twibbon.name if twibbon else None
    if isinstance(input_files, str) and input_files.lower().endswith(('.zip', '.rar')):
        return process_images(
            input_files, bg_method, watermark_path, twibbon_path,
            canvas_size, output_format, bg_choice, custom_color, num_workers,
            rotation, direction, flip, True, progress, sheet_file, use_qwen, 
            snap_to_bottom, snap_to_top, snap_to_left, snap_to_right, auto_snap
        )
    elif isinstance(input_files, list):
        return process_images(
            input_files, bg_method, watermark_path, twibbon_path,
            canvas_size, output_format, bg_choice, custom_color, num_workers,
            rotation, direction, flip, True, progress, sheet_file, use_qwen, 
            snap_to_bottom, snap_to_top, snap_to_left, snap_to_right, auto_snap
        )
    else:
        return process_images(
            input_files.name, bg_method, watermark_path, twibbon_path,
            canvas_size, output_format, bg_choice, custom_color, num_workers,
            rotation, direction, flip, True, progress, sheet_file, use_qwen, 
            snap_to_bottom, snap_to_top, snap_to_left, snap_to_right, auto_snap
        )

def show_color_picker(bg_choice):
    if bg_choice == 'custom':
        return gr.update(visible=True)
    return gr.update(visible=False)

def show_custom_canvas(canvas_size):
    if canvas_size == 'Custom':
        return gr.update(visible=True), gr.update(visible=True)
    return gr.update(visible=False), gr.update(visible=False)

def parse_color(color_str):
    """Convert color string to format that PIL can understand"""
    if not color_str:
        return "#ffffff"
    
    # If it's already a hex color, return as-is
    if color_str.startswith('#'):
        return color_str
    
    # Handle rgba() format from Gradio ColorPicker
    if color_str.startswith('rgba(') or color_str.startswith('rgb('):
        import re
        # Extract numbers from rgba(r, g, b, a) or rgb(r, g, b)
        numbers = re.findall(r'[\d.]+', color_str)
        if len(numbers) >= 3:
            r = int(float(numbers[0]))
            g = int(float(numbers[1]))
            b = int(float(numbers[2]))
            # Convert to hex
            return f"#{r:02x}{g:02x}{b:02x}"
    
    # Default fallback
    return "#ffffff"

def update_compare(evt: gr.SelectData, classifications):
    if isinstance(evt.value, dict) and 'caption' in evt.value:
        in_path = evt.value['caption'].split("Input: ")[-1]
        out_path = evt.value['image']['path']
        orig = Image.open(in_path)
        proc = Image.open(out_path)
        ratio_o = f"{orig.width}x{orig.height}"
        ratio_p = f"{proc.width}x{proc.height}"
        filename = os.path.basename(in_path)
        if filename in classifications:
            cls = classifications[filename]["classification"]
            pad = classifications[filename]["padding"]
            selected_info_text = f"Classification: {cls}, Padding - Top: {pad['top']}, Bottom: {pad['bottom']}, Left: {pad['left']}, Right: {pad['right']}"
        else:
            selected_info_text = "No classification data available"
        return (
            gr.update(value=in_path),
            gr.update(value=out_path),
            gr.update(value=ratio_o),
            gr.update(value=ratio_p),
            gr.update(value=selected_info_text)
        )
    else:
        print("No caption found in selection.")
        return (
            gr.update(value=None),
            gr.update(value=None),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value="Select an image to see details")
        )

def process(
    input_files,
    bg_method,
    watermark,
    twibbon,
    canvas_size,
    output_format,
    bg_choice,
    custom_color,
    num_workers,
    rotation=None,
    direction=None,
    flip=False,
    sheet_file=None,
    use_qwen_str="Default (No Vision)",
    snap_to_bottom=False,
    snap_to_top=False,
    snap_to_left=False,
    snap_to_right=False,
    auto_snap=False,
    canvas_width=1080,
    canvas_height=1080
):
    use_qwen = (use_qwen_str == "Utilize Vision Model")  # Convert string to boolean
    
    # Handle custom canvas size
    if canvas_size == 'Custom':
        canvas_size = (canvas_width, canvas_height)
    
    _, procd, zip_out, tt, classifications = gradio_interface(
        input_files, bg_method, watermark, twibbon,
        canvas_size, output_format, bg_choice, custom_color, num_workers,
        rotation, direction, flip, sheet_file, use_qwen, snap_to_bottom, snap_to_top, snap_to_left, snap_to_right, auto_snap
        )
    if not procd:
        return [], None, "No Image Processed.", "No Classification Available", {}
    result_g = []
    for outf, inf in procd:
        if not os.path.exists(outf):
            print(f"[ERROR] Missing out: {outf}")
            continue
        result_g.append((outf, f"Input: {inf}"))
    class_text = "\n".join([
        f"{img}: Classification - {data['classification']}, Padding - Top: {data['padding']['top']}, Bottom: {data['padding']['bottom']}, Left: {data['padding']['left']}, Right: {data['padding']['right']}"
        for img, data in classifications.items()
    ]) or "No classifications recorded."
    return result_g, zip_out, f"{tt:.2f} seconds", class_text, classifications

def stop_processing():
    stop_event.set()

def preset_snap_rules(filename, image_path=None):
    """
    Menerapkan aturan preset untuk snap settings berdasarkan nama file atau kategori
    Returns dict dengan format {'snap_top': bool, 'snap_right': bool, 'snap_bottom': bool, 'snap_left': bool}
    """
    filename_lower = filename.lower()
    
    # Default settings
    settings = {
        'snap_top': False,
        'snap_right': False,
        'snap_bottom': False,
        'snap_left': False
    }
    
    # ---- Pola untuk produk berdasarkan urutan gambar ----
    # Angka di filename biasanya menunjukkan view produk
    view_num = None
    for pattern in ['_01', '_02', '_03', '_04', '_05', '_06', '_1.', '_2.', '_3.', '_4.', '_5.', '_6.']:
        if pattern in filename_lower:
            view_num = int(pattern.strip('_.'))
            break
    
    # --- Pola Format Pendek (pakaian renang, baju, pakaian olahraga) ---
    # Format: @1000xxxxxx_01.jpg, @1000xxxxxx_02.jpg, dll
    if filename_lower.startswith('@10002'):
        print(f"Matched special pattern @10002xxxxx for {filename}")
        # View pertama biasanya depan, snap_bottom
        if view_num == 1:
            settings['snap_bottom'] = True
            settings['snap_left'] = True
        # View kedua biasanya belakang, snap_bottom
        elif view_num == 2:
            settings['snap_bottom'] = True
            settings['snap_right'] = True
        # View ketiga biasanya samping, snap_left dan snap_bottom
        elif view_num == 3:
            settings['snap_bottom'] = True
            settings['snap_left'] = True
            settings['snap_top'] = True
        # View keempat biasanya samping lain, snap_right dan snap_bottom
        elif view_num == 4:
            settings['snap_bottom'] = True
            settings['snap_right'] = True
            settings['snap_top'] = True
    
    # --- Pola Bikini/Baju Renang ---
    elif any(x in filename_lower for x in ['bikini', 'swimwear', 'swimsuit', 'swim']):
        # Untuk bikini tops (hanya bagian atas)
        if any(x in filename_lower for x in ['top', 'bra', 'bust']):
            if view_num == 1:  # Foto produk pertama - biasanya depan
                settings['snap_bottom'] = True
            elif view_num == 2:  # Foto produk kedua - biasanya belakang
                settings['snap_bottom'] = True
            elif view_num == 3:  # Foto produk ketiga - biasanya samping
                settings['snap_bottom'] = True
                settings['snap_left'] = True
            elif view_num == 4:  # Foto produk keempat - biasanya samping lain
                settings['snap_bottom'] = True
                settings['snap_right'] = True
        # Untuk bikini bottoms (hanya bagian bawah)
        elif any(x in filename_lower for x in ['bottom', 'pant', 'brief']):
            if view_num == 1:  # Foto produk pertama - biasanya depan
                settings['snap_bottom'] = True
            elif view_num == 2:  # Foto produk kedua - biasanya belakang
                settings['snap_bottom'] = True
            elif view_num == 3:  # Foto produk ketiga - biasanya samping
                settings['snap_bottom'] = True
                settings['snap_left'] = True
                settings['snap_top'] = True
            elif view_num == 4:  # Foto produk keempat - biasanya samping lain
                settings['snap_bottom'] = True
                settings['snap_right'] = True
                settings['snap_top'] = True
        # Untuk one-piece atau bikini sets
        else:
            if view_num == 1:  # Foto produk pertama - biasanya depan
                settings['snap_bottom'] = True
            elif view_num == 2:  # Foto produk kedua - biasanya belakang
                settings['snap_bottom'] = True
            elif view_num == 3:  # Foto produk ketiga - biasanya samping
                settings['snap_bottom'] = True
                settings['snap_left'] = True
            elif view_num == 4:  # Foto produk keempat - biasanya samping lain
                settings['snap_bottom'] = True
                settings['snap_right'] = True
    
    # --- Pola Pakaian Dengan Model ---
    elif any(x in filename_lower for x in ['_model_', 'human', 'person']):
        settings['snap_bottom'] = True
        # Jika terlihat dari samping, tambahkan snap kiri atau kanan
        if "_left" in filename_lower or "_samping" in filename_lower:
            settings['snap_left'] = True
        if "_right" in filename_lower:
            settings['snap_right'] = True
    
    # --- Pola untuk Tas ---
    elif any(x in filename_lower for x in ['bag', 'backpack', 'tas', 'sling']):
        # Format kode file tertentu
        if view_num == 1:  # View depan
            settings['snap_bottom'] = True
        elif view_num == 2:  # View belakang 
            settings['snap_bottom'] = True
        elif view_num == 3:  # View samping
            settings['snap_bottom'] = True
            settings['snap_left'] = True
        elif view_num == 4:  # View samping lain
            settings['snap_bottom'] = True
            settings['snap_right'] = True
    
    # --- Pola untuk Sepatu ---
    elif any(x in filename_lower for x in ['shoe', 'footwear', 'sepatu']):
        if "_side" in filename_lower or "_samping" in filename_lower:
            settings['snap_bottom'] = True
            if "_left" in filename_lower:
                settings['snap_left'] = True
            elif "_right" in filename_lower:
                settings['snap_right'] = True
            else:
                # Default untuk sepatu dari samping (biasanya sepatu kiri)
                settings['snap_left'] = True
    
    # --- Kasus khusus berdasarkan nama file persis ---
    # Contoh file yang disebutkan user
    if "1000218277_01" in filename_lower:
        settings['snap_bottom'] = True
        settings['snap_left'] = True
    elif "1000218265_01" in filename_lower:
        settings['snap_top'] = True
        settings['snap_bottom'] = True
        settings['snap_left'] = True
    elif "1000218268_01" in filename_lower:
        settings['snap_top'] = True
        settings['snap_bottom'] = True
        settings['snap_right'] = True
    
    # Kasus khusus untuk pola @1000xxxxxx (seperti yang disebutkan user)
    elif filename_lower.startswith('@'):
        if '_01' in filename_lower and filename_lower.startswith('@10002'):
            settings['snap_bottom'] = True 
            settings['snap_left'] = True
    
    # Tambahkan lebih banyak pola sesuai kebutuhan
    
    return settings

with gr.Blocks(theme='allenai/gradio-theme') as iface:
    gr.Markdown("## Image BG Removal with Rotation, Watermark, Twibbon & Classifications for Padding Override")
    with gr.Row():
        input_files = gr.File(label="Upload (Image(s)/ZIP/RAR)", file_types=[".zip", ".rar", "image"], interactive=True)
        watermark = gr.File(label="Watermark (Optional)", file_types=[".png"])
        twibbon = gr.File(label="Twibbon (Optional)", file_types=[".png"])
        sheet_file = gr.File(label="Upload Sheet (.xlsx/.csv)", file_types=[".xlsx", ".csv"], interactive=True)
    with gr.Row():
        bg_method = gr.Radio(["bria", "none"],
                            label="Background Removal", value="bria")
        bg_choice = gr.Radio(["transparent", "white", "custom"], label="BG Choice", value="white")
        custom_color = gr.ColorPicker(label="Custom BG", value="#ffffff", visible=False)
        output_format = gr.Radio(["PNG", "JPG"], label="Output Format", value="JPG")
        num_workers = gr.Slider(1, 16, 1, label="Number of Workers", value=5)
        use_qwen = gr.Dropdown(
            ["Default (No Vision)", "Utilize Vision Model"],
            label="Classification",
            value="Default (No Vision)"  # Default is off
        )
    with gr.Row():
        canvas_size = gr.Radio(
            choices=[
                "primer-sale.psd", "Custom"
            ],
            label="Canvas Size", value="primer-sale.psd"
        )
    with gr.Row() as custom_canvas_row:
        canvas_width = gr.Number(label="Canvas Width (px)", value=1080, minimum=1, maximum=5000, step=1, visible=False)
        canvas_height = gr.Number(label="Canvas Height (px)", value=1080, minimum=1, maximum=5000, step=1, visible=False)
    with gr.Row():
        rotation = gr.Radio(["None", "90 Degrees", "180 Degrees"], label="Rotation Angle", value="None")
        direction = gr.Radio(["None", "Clockwise", "Anticlockwise"], label="Direction", value="None")
        flip_option = gr.Checkbox(label="Flip Horizontal", value=False)
        auto_snap = gr.Checkbox(label="Auto Snap (Gunakan AI untuk menentukan snap setting)", value=False)
    
    # Kelompokkan semua snap manual di baris yang terpisah
    with gr.Row() as manual_snap_row:
        gr.Markdown("### Manual Snap Settings (tidak digunakan jika Auto Snap aktif)")
        snap_to_bottom = gr.Checkbox(label="Snap to Bottom (Force padding bottom 0)", value=False)
        snap_to_top = gr.Checkbox(label="Snap to Top (Force padding top 0)", value=False)
        snap_to_left = gr.Checkbox(label="Snap to Left (Force padding left 0)", value=False)
        snap_to_right = gr.Checkbox(label="Snap to Right (Force padding right 0)", value=False)
        
    proc_btn = gr.Button("Process Images")
    stop_btn = gr.Button("Stop")
    with gr.Row():
        gallery_processed = gr.Gallery(label="Processed Images")
    with gr.Row():
        selected_info = gr.Textbox(label="Selected Image Classification and Padding", lines=2, interactive=False)
    with gr.Row():
        img_orig = gr.Image(label="Original", interactive=False)
        img_proc = gr.Image(label="Processed", interactive=False)
    with gr.Row():
        ratio_orig = gr.Textbox(label="Original Ratio")
        ratio_proc = gr.Textbox(label="Processed Ratio")
    with gr.Row():
        out_zip = gr.File(label="Download as ZIP")
        time_box = gr.Textbox(label="Processing Time (seconds)")
        classifications_state = gr.State()
    with gr.Row():
        class_display = gr.Textbox(label="All Classification and Padding Results", lines=5, interactive=False)

    bg_choice.change(show_color_picker, inputs=bg_choice, outputs=custom_color)
    canvas_size.change(show_custom_canvas, inputs=canvas_size, outputs=[canvas_width, canvas_height])
    proc_btn.click(
        fn=process,
        inputs=[input_files, bg_method, watermark, twibbon, canvas_size, output_format,
                bg_choice, custom_color, num_workers, rotation, direction, flip_option, 
                sheet_file, use_qwen, snap_to_bottom, snap_to_top, snap_to_left, snap_to_right, 
                auto_snap, canvas_width, canvas_height],
        outputs=[gallery_processed, out_zip, time_box, class_display, classifications_state]
    )
    gallery_processed.select(
        update_compare,
        inputs=[classifications_state],
        outputs=[img_orig, img_proc, ratio_orig, ratio_proc, selected_info]
    )
    stop_btn.click(fn=stop_processing, outputs=[])

    # Add dependency for hiding/showing manual snap options
    def update_manual_snap_visibility(auto_snap_active):
        return gr.update(visible=not auto_snap_active)
        
    auto_snap.change(
        fn=update_manual_snap_visibility,
        inputs=[auto_snap],
        outputs=[manual_snap_row]
    )

iface.launch(share=True)