src/streamlit_app.py

import streamlit as st
from PIL import Image, ImageColor, ImageDraw, ImageFont, PngImagePlugin
import torch
import torch.nn.functional as F
from torchvision import transforms
from transformers import AutoModelForImageSegmentation, AutoImageProcessor, Swin2SRForImageSuperResolution, VitMatteForImageMatting
import io
import numpy as np
import gc

# Page Configuration
st.set_page_config(layout="wide", page_title="AI Image Lab Pro")

# --- 1. MODEL LOADING (Cached - UNCHANGED) ---

@st.cache_resource
def load_rmbg_model():
    model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, device

@st.cache_resource
def load_birefnet_model():
    model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, device

@st.cache_resource
def load_vitmatte_model():
    processor = AutoImageProcessor.from_pretrained("hustvl/vitmatte-small-composition-1k")
    model = VitMatteForImageMatting.from_pretrained("hustvl/vitmatte-small-composition-1k")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return processor, model, device

@st.cache_resource
def load_upscaler(scale=2):
    if scale == 4:
        model_id = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr"
    else:
        model_id = "caidas/swin2SR-classical-sr-x2-64"
    processor = AutoImageProcessor.from_pretrained(model_id)
    model = Swin2SRForImageSuperResolution.from_pretrained(model_id)
    return processor, model

# --- 2. HELPER FUNCTIONS (AI & Processing - UNCHANGED) ---

def cleanup_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def find_mask_tensor(output):
    if isinstance(output, torch.Tensor):
        if output.dim() == 4 and output.shape[1] == 1: return output
        elif output.dim() == 3 and output.shape[0] == 1: return output
        return None
    if hasattr(output, "logits"): return find_mask_tensor(output.logits)
    elif isinstance(output, (list, tuple)):
        for item in output:
            found = find_mask_tensor(item)
            if found is not None: return found
    return None

def generate_trimap(mask_tensor, erode_kernel_size=10, dilate_kernel_size=10):
    if mask_tensor.dim() == 3: mask_tensor = mask_tensor.unsqueeze(0)
    erode_k = erode_kernel_size
    dilate_k = dilate_kernel_size
    dilated = F.max_pool2d(mask_tensor, kernel_size=dilate_k, stride=1, padding=dilate_k//2)
    eroded = -F.max_pool2d(-mask_tensor, kernel_size=erode_k, stride=1, padding=erode_k//2)
    trimap = torch.full_like(mask_tensor, 0.5)
    trimap[eroded > 0.5] = 1.0
    trimap[dilated < 0.5] = 0.0
    return trimap

# --- 3. INFERENCE LOGIC (UNCHANGED) ---

def inference_segmentation(model, image, device, resolution=1024):
    w, h = image.size
    transform = transforms.Compose([
        transforms.Resize((resolution, resolution)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    input_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(input_tensor)
    
    result_tensor = find_mask_tensor(outputs)
    if result_tensor is None: result_tensor = outputs[0] if isinstance(outputs, (list, tuple)) else outputs
    if not isinstance(result_tensor, torch.Tensor):
         if isinstance(result_tensor, (list, tuple)): result_tensor = result_tensor[0]

    pred = result_tensor.squeeze().cpu()
    if pred.max() > 1 or pred.min() < 0: pred = pred.sigmoid()
    
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize((w, h), resample=Image.LANCZOS)
    return mask

def inference_vitmatte(image, device):
    cleanup_memory()
    original_size = image.size
    max_dim = 1536
    if max(image.size) > max_dim:
        scale_ratio = max_dim / max(image.size)
        new_w = int(image.size[0] * scale_ratio)
        new_h = int(image.size[1] * scale_ratio)
        processing_image = image.resize((new_w, new_h), Image.LANCZOS)
    else:
        processing_image = image

    rmbg_model, _ = load_rmbg_model() 
    rough_mask_pil = inference_segmentation(rmbg_model, processing_image, device, resolution=1024)
    
    mask_tensor = transforms.ToTensor()(rough_mask_pil).to(device)
    trimap_tensor = generate_trimap(mask_tensor, erode_kernel_size=25, dilate_kernel_size=25)
    trimap_pil = transforms.ToPILImage()(trimap_tensor.squeeze().cpu())
    
    processor, model, _ = load_vitmatte_model()
    inputs = processor(images=processing_image, trimaps=trimap_pil, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    alphas = outputs.alphas
    alpha_np = alphas.squeeze().cpu().numpy()
    alpha_pil = Image.fromarray((alpha_np * 255).astype("uint8"), mode="L")
    
    if original_size != processing_image.size:
        alpha_pil = alpha_pil.resize(original_size, resample=Image.LANCZOS)
    
    cleanup_memory()
    return alpha_pil

@st.cache_data(show_spinner=False)
def process_background_removal(image_bytes, method="RMBG-1.4"):
    cleanup_memory()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGBA")
    image_rgb = image.convert("RGB")
    
    if method == "RMBG-1.4":
        model, device = load_rmbg_model()
        mask = inference_segmentation(model, image_rgb, device)
        
    elif method == "BiRefNet (Heavy)":
        model, device = load_birefnet_model()
        mask = inference_segmentation(model, image_rgb, device, resolution=1024)
        
    elif method == "VitMatte (Refiner)":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        mask = inference_vitmatte(image_rgb, device)
    
    else:
        return image

    final_image = image_rgb.copy()
    final_image.putalpha(mask)
    return final_image

# --- Upscaling Logic ---
def run_swin_inference(image, processor, model):
    inputs = processor(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
    output = np.moveaxis(output, 0, -1)
    output = (output * 255.0).round().astype(np.uint8)
    return Image.fromarray(output)

def upscale_chunk_logic(image, processor, model):
    if image.mode == 'RGBA':
        r, g, b, a = image.split()
        rgb_image = Image.merge('RGB', (r, g, b))
        upscaled_rgb = run_swin_inference(rgb_image, processor, model)
        upscaled_a = a.resize(upscaled_rgb.size, Image.Resampling.LANCZOS)
        return Image.merge('RGBA', (*upscaled_rgb.split(), upscaled_a))
    else:
        return run_swin_inference(image, processor, model)

def process_tiled_upscale(image, scale_factor, grid_n, progress_bar):
    cleanup_memory()
    processor, model = load_upscaler(scale_factor)
    w, h = image.size
    rows = cols = grid_n
    tile_w = w // cols
    tile_h = h // rows
    overlap = 32 
    full_image = Image.new(image.mode, (w * scale_factor, h * scale_factor))
    total_tiles = rows * cols
    count = 0
    for y in range(rows):
        for x in range(cols):
            target_left = x * tile_w
            target_upper = y * tile_h
            target_right = w if x == cols - 1 else (x + 1) * tile_w
            target_lower = h if y == rows - 1 else (y + 1) * tile_h
            source_left = max(0, target_left - overlap)
            source_upper = max(0, target_upper - overlap)
            source_right = min(w, target_right + overlap)
            source_lower = min(h, target_lower + overlap)
            tile = image.crop((source_left, source_upper, source_right, source_lower))
            upscaled_tile = upscale_chunk_logic(tile, processor, model)
            target_w = target_right - target_left
            target_h = target_lower - target_upper
            extra_left = target_left - source_left
            extra_upper = target_upper - source_upper
            crop_x = extra_left * scale_factor
            crop_y = extra_upper * scale_factor
            crop_w = target_w * scale_factor
            crop_h = target_h * scale_factor
            clean_tile = upscaled_tile.crop((crop_x, crop_y, crop_x + crop_w, crop_y + crop_h))
            paste_x = target_left * scale_factor
            paste_y = target_upper * scale_factor
            full_image.paste(clean_tile, (paste_x, paste_y))
            del tile, upscaled_tile, clean_tile
            cleanup_memory()
            count += 1
            progress_bar.progress(count / total_tiles, text=f"Upscaling Tile {count}/{total_tiles}...")
    return full_image

# --- 4. NEW HELPER FUNCTIONS (Watermark & Metadata) ---

def apply_watermark(image, text, opacity, size_scale, position):
    if not text: return image
    watermark_image = image.convert("RGBA")
    text_layer = Image.new("RGBA", watermark_image.size, (255, 255, 255, 0))
    draw = ImageDraw.Draw(text_layer)
    w, h = watermark_image.size
    base_font_size = int(h * 0.05)
    font_size = int(base_font_size * size_scale)
    try:
        font = ImageFont.load_default()
    except ImportError:
         font = ImageFont.load_default()
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    padding = 20
    x, y = 0, 0
    if position == "Bottom Right":
        x, y = w - text_width - padding, h - text_height - padding
    elif position == "Bottom Left":
        x, y = padding, h - text_height - padding
    elif position == "Top Right":
        x, y = w - text_width - padding, padding
    elif position == "Top Left":
        x, y = padding, padding
    elif position == "Center":
        x, y = (w - text_width) // 2, (h - text_height) // 2
    alpha_val = int(opacity * 255)
    text_color = (255, 255, 255, alpha_val)
    draw.text((x, y), text, font=font, fill=text_color)
    output = Image.alpha_composite(watermark_image, text_layer)
    if image.mode == 'RGB': return output.convert('RGB')
    return output

def convert_image_to_bytes_with_metadata(img, author=None, copyright_text=None):
    buf = io.BytesIO()
    pnginfo = PngImagePlugin.PngInfo()
    if author:
        pnginfo.add_text("Author", author)
        pnginfo.add_text("Software", "AI Image Lab Pro")
    if copyright_text:
        pnginfo.add_text("Copyright", copyright_text)
    img.save(buf, format="PNG", pnginfo=pnginfo)
    return buf.getvalue()

# --- 5. MAIN APP ---

def main():
    st.title("✨ AI Image Lab: Professional")

    # --- Sidebar Section 1: Input & Metadata ---
    st.sidebar.header("1. Input & Metadata")
    uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "webp"])
    
    clean_metadata_on_load = st.sidebar.checkbox("Strip Original Metadata on Load", value=False)

    if uploaded_file is not None:
        file_bytes = uploaded_file.getvalue()
        initial_img_inspect = Image.open(io.BytesIO(file_bytes))
        with st.sidebar.expander("🔍 View Original Metadata"):
            if initial_img_inspect.info:
                safe_info = {k: v for k, v in initial_img_inspect.info.items() if isinstance(v, (str, int, float))}
                if safe_info: st.json(safe_info)
                else: st.write("Binary metadata hidden.")
            else: st.write("No metadata found.")

        if clean_metadata_on_load:
            clean_img = Image.new(initial_img_inspect.mode, initial_img_inspect.size)
            clean_img.putdata(list(initial_img_inspect.getdata()))
            buf = io.BytesIO()
            clean_img.save(buf, format="PNG")
            processing_bytes = buf.getvalue()
            st.sidebar.success("Metadata stripped.")
        else:
            processing_bytes = file_bytes

    # --- Sidebar Section 2: AI Processing ---
    st.sidebar.header("2. AI Processing")
    remove_bg = st.sidebar.checkbox("Remove Background", value=True)
    
    if remove_bg:
        bg_model = st.sidebar.selectbox("AI Model", ["BiRefNet (Heavy)", "RMBG-1.4", "VitMatte (Refiner)"], index=0)
    else:
        bg_model = "None"

    upscale_mode = st.sidebar.radio("Magnification", ["None", "2x", "4x"])
    if upscale_mode != "None":
        grid_n = st.sidebar.slider("Grid Split", 2, 8, 4)
    else:
        grid_n = 2

    # --- Sidebar Section 3: Studio Tools ---
    st.sidebar.markdown("---")
    st.sidebar.header("3. Studio Tools")
    
    bg_color_mode = st.sidebar.selectbox("Background Color", ["Transparent", "White", "Black", "Custom"])
    custom_bg_color = "#FFFFFF"
    if bg_color_mode == "Custom":
        custom_bg_color = st.sidebar.color_picker("Pick color", "#FF0000")

    enable_smart_crop = st.sidebar.checkbox("Smart Auto-Crop (to Subject)", value=False)
    crop_padding = 0
    if enable_smart_crop:
        crop_padding = st.sidebar.slider("Auto-Crop Padding", 0, 500, 50)
    
    st.sidebar.caption("Manual Crop (px)")
    col_c1, col_c2 = st.sidebar.columns(2)
    with col_c1:
        crop_top = st.number_input("Top", min_value=0, value=0, step=10)
        crop_left = st.number_input("Left", min_value=0, value=0, step=10)
    with col_c2:
        crop_bottom = st.number_input("Bottom", min_value=0, value=0, step=10)
        crop_right = st.number_input("Right", min_value=0, value=0, step=10)

    rotate_angle = st.sidebar.slider("Rotate", -180, 180, 0, 1)

    st.sidebar.subheader("Watermark")
    wm_text = st.sidebar.text_input("Watermark Text")
    wm_opacity = st.sidebar.slider("Opacity", 0.1, 1.0, 0.5)
    wm_size = st.sidebar.slider("Size Scale", 0.5, 3.0, 1.0)
    wm_position = st.sidebar.selectbox("Position", ["Bottom Right", "Bottom Left", "Top Right", "Top Left", "Center"])


    # --- Sidebar Section 4: Output Settings ---
    st.sidebar.markdown("---")
    st.sidebar.header("4. Output Settings")
    meta_author = st.sidebar.text_input("Author Name")
    meta_copyright = st.sidebar.text_input("Copyright Notice")


    # --- Main Application Logic ---
    if uploaded_file is not None:
        if remove_bg:
            with st.spinner(f"Removing background using {bg_model}..."):
                processed_image = process_background_removal(processing_bytes, bg_model)
        else:
            processed_image = Image.open(io.BytesIO(processing_bytes)).convert("RGBA")

        if upscale_mode != "None":
            scale = 4 if "4x" in upscale_mode else 2
            cache_key = f"{uploaded_file.name}_clean{clean_metadata_on_load}_{bg_model}_{scale}_{grid_n}_v11"
            if "upscale_cache" not in st.session_state: st.session_state.upscale_cache = {}
            if cache_key in st.session_state.upscale_cache:
                processed_image = st.session_state.upscale_cache[cache_key]
                st.info("✅ Loaded upscaled image from cache")
            else:
                progress_bar = st.progress(0, text="Initializing AI models...")
                processed_image = process_tiled_upscale(processed_image, scale, grid_n, progress_bar)
                progress_bar.empty()
                st.session_state.upscale_cache[cache_key] = processed_image
        
        final_image = processed_image.copy()

        # A. Rotation
        if rotate_angle != 0:
            final_image = final_image.rotate(rotate_angle, expand=True)

        # B. Smart Auto-Crop
        if enable_smart_crop and final_image.mode == 'RGBA':
            alpha = final_image.getchannel('A')
            bbox = alpha.getbbox()
            if bbox:
                left, upper, right, lower = bbox
                w, h = final_image.size
                left = max(0, left - crop_padding)
                upper = max(0, upper - crop_padding)
                right = min(w, right + crop_padding)
                lower = min(h, lower + crop_padding)
                final_image = final_image.crop((left, upper, right, lower))

        # C. Manual Crop
        # Applied after Smart Crop so you can refine it
        w, h = final_image.size
        # Ensure we don't crop beyond image dimensions
        valid_left = min(crop_left, w - 1)
        valid_top = min(crop_top, h - 1)
        valid_right = min(crop_right, w - valid_left - 1)
        valid_bottom = min(crop_bottom, h - valid_top - 1)
        
        if valid_left > 0 or valid_top > 0 or valid_right > 0 or valid_bottom > 0:
            final_image = final_image.crop((
                valid_left, 
                valid_top, 
                w - valid_right, 
                h - valid_bottom
            ))

        # D. Background Compositing
        if bg_color_mode != "Transparent" and final_image.mode == 'RGBA':
            if bg_color_mode == "White": bg = Image.new("RGBA", final_image.size, "WHITE")
            elif bg_color_mode == "Black": bg = Image.new("RGBA", final_image.size, "BLACK")
            else: bg = Image.new("RGBA", final_image.size, custom_bg_color)
            bg.alpha_composite(final_image)
            final_image = bg.convert("RGB")

        # E. Watermark
        if wm_text:
            final_image = apply_watermark(final_image, wm_text, wm_opacity, wm_size, wm_position)

        # --- Display ---
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("Original")
            st.image(Image.open(io.BytesIO(file_bytes)), use_container_width=True)
        
        with col2:
            st.subheader("Result")
            st.markdown("""<style>[data-testid="stImage"] {background-image: url('https://i.imgur.com/s1B49hR.png'); background-size: 20px 20px;}</style>""", unsafe_allow_html=True)
            st.image(final_image, use_container_width=True)

        st.markdown("---")
        download_data = convert_image_to_bytes_with_metadata(final_image, author=meta_author, copyright_text=meta_copyright)
        st.download_button(
            label="💾 Download Result (PNG with Metadata)",
            data=download_data,
            file_name="processed_image.png",
            mime="image/png"
        )

if __name__ == "__main__":
    main()