Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,447 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 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"
+        )
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if __name__ == "__main__":
+    main()