import streamlit as st from PIL import Image, ImageEnhance import torch from torchvision import transforms from transformers import AutoModelForImageSegmentation, AutoImageProcessor, Swin2SRForImageSuperResolution import io import numpy as np # Page Configuration st.set_page_config(layout="wide", page_title="AI Image Lab") # --- 1. MODEL LOADING (Cached) --- @st.cache_resource def load_rembg_model(): """Loads RMBG-1.4 for Background Removal.""" # We use 'briaai/RMBG-1.4' 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_upscaler(scale=2): """Loads Swin2SR for Super-Resolution (2x or 4x).""" if scale == 4: model_id = "caidas/swin2SR-classical-sr-x4-63" 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. PROCESSING FUNCTIONS --- def find_mask_tensor(output): """ Recursively searches any nested structure (list, tuple, dict, object) to find the first Tensor that looks like a mask (1 channel). """ # 1. If it's a Tensor, check if it's the mask we want if isinstance(output, torch.Tensor): # We look for shape [Batch, 1, H, W] or [1, H, W] # It must have 1 channel (index 1 for 4D, index 0 for 3D) if output.dim() == 4 and output.shape[1] == 1: return output elif output.dim() == 3 and output.shape[0] == 1: return output # If it has > 1 channels (e.g. 64), it's a feature map, ignore it. return None # 2. If it's a Dict/ModelOutput (like .logits), check values if hasattr(output, "items"): for val in output.values(): found = find_mask_tensor(val) if found is not None: return found # Special case for Hugging Face model outputs with attributes elif hasattr(output, "logits"): return find_mask_tensor(output.logits) # 3. If it's a List or Tuple, iterate through elements elif isinstance(output, (list, tuple)): for item in output: found = find_mask_tensor(item) if found is not None: return found return None def safe_rembg_inference(model, image, device): """ Robust inference for RMBG-1.4 using Deep Search. """ w, h = image.size # Preprocessing transform_image = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) input_images = transform_image(image).unsqueeze(0).to(device) # Inference with torch.no_grad(): outputs = model(input_images) # --- DEEP SEARCH FOR MASK --- result_tensor = find_mask_tensor(outputs) if result_tensor is None: # Fallback: If deep search failed, try just grabbing the first tensor found # (Even if dimensions look weird, it's better than crashing) if isinstance(outputs, (list, tuple)): result_tensor = outputs[0] else: result_tensor = outputs # Post-processing # Ensure it's a tensor before operations if not isinstance(result_tensor, torch.Tensor): # If we still have a list here, we take the first element blindly if isinstance(result_tensor, (list, tuple)): result_tensor = result_tensor[0] pred = result_tensor.squeeze().cpu() # Sometimes output is already sigmoid, sometimes logits. # If values are > 1 or < 0, apply sigmoid. if pred.max() > 1 or pred.min() < 0: pred = pred.sigmoid() # Convert mask to PIL pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize((w, h)) # Apply mask image.putalpha(mask) return image def ai_upscale(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 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 convert_image_to_bytes(img): buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() # --- 3. MAIN APP --- def main(): st.title("✨ AI Image Lab: Robust Edition") st.markdown("Features: **RMBG-1.4 (Pure PyTorch)** | **Swin2SR (Upscaling)** | **Geometry**") # --- Sidebar --- st.sidebar.header("1. Background") remove_bg = st.sidebar.checkbox("Remove Background", value=False) st.sidebar.header("2. AI Upscaling") upscale_mode = st.sidebar.radio("Magnification", ["None", "2x (Fast)", "4x (Slow)"]) st.sidebar.header("3. Geometry") rotate_angle = st.sidebar.slider("Rotate", -180, 180, 0, 1) # --- Main --- uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "webp"]) if uploaded_file is not None: image = Image.open(uploaded_file).convert("RGB") processed_image = image.copy() # 1. Background if remove_bg: st.info("Loading RMBG Model...") try: bg_model, device = load_rembg_model() with st.spinner("Removing background..."): processed_image = safe_rembg_inference(bg_model, processed_image, device) except Exception as e: st.error(f"Background Removal Failed: {e}") # 2. Upscaling if upscale_mode != "None": scale = 4 if "4x" in upscale_mode else 2 st.info(f"Loading Swin2SR x{scale} Model...") try: processor, upscaler = load_upscaler(scale) with st.spinner(f"Upscaling x{scale}..."): processed_image = ai_upscale(processed_image, processor, upscaler) except Exception as e: st.error(f"Upscaling Failed: {e}") # 3. Rotation if rotate_angle != 0: processed_image = processed_image.rotate(rotate_angle, expand=True) # --- Display --- col1, col2 = st.columns(2) with col1: st.subheader("Original") st.image(image, use_container_width=True) st.caption(f"Size: {image.size}") with col2: st.subheader("Result") st.image(processed_image, use_container_width=True) st.caption(f"Size: {processed_image.size}") # --- Download --- st.markdown("---") st.download_button( label="💾 Download Result (PNG)", data=convert_image_to_bytes(processed_image), file_name="processed_image.png", mime="image/png" ) if __name__ == "__main__": main()