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("""""", 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()