app.py

import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import SamModel, SamProcessor
from diffusers import StableDiffusionInpaintPipeline
import requests
from io import BytesIO

# Set up device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load SAM model for segmentation
print("Loading SAM model...")
sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

# Load Stable Diffusion for outpainting
print("Loading Stable Diffusion model...")
inpaint_model = StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
).to(device)

def get_sam_mask(image, points=None):
    """Get segmentation mask using SAM model"""
    if points is None:
        # If no points provided, use center point
        height, width = image.shape[:2]
        points = [[[width // 2, height // 2]]]
        
    # Convert to PIL if needed
    if not isinstance(image, Image.Image):
        image_pil = Image.fromarray(image)
    else:
        image_pil = image
        
    # Process the image and point prompts
    inputs = sam_processor(
        images=image_pil,
        input_points=points,
        return_tensors="pt"
    ).to(device)
    
    # Generate mask
    with torch.no_grad():
        outputs = sam_model(**inputs)
        masks = sam_processor.image_processor.post_process_masks(
            outputs.pred_masks.cpu(),
            inputs["original_sizes"].cpu(),
            inputs["reshaped_input_sizes"].cpu()
        )
    
    # Get the mask
    mask = masks[0][0].numpy()
    return mask

def adjust_aspect_ratio(image, mask, target_ratio, prompt=""):
    """Adjust image to target aspect ratio while preserving important content"""
    # Convert PIL to numpy if needed
    if isinstance(image, Image.Image):
        image_np = np.array(image)
    else:
        image_np = image
        
    h, w = image_np.shape[:2]
    current_ratio = w / h
    target_ratio_value = eval(target_ratio.replace(':', '/'))
    
    # Determine if we need to add width or height
    if current_ratio < target_ratio_value:
        # Need to add width (outpaint left/right)
        new_width = int(h * target_ratio_value)
        new_height = h
        
        # Calculate padding
        pad_width = new_width - w
        pad_left = pad_width // 2
        pad_right = pad_width - pad_left
        
        # Create canvas with padding
        result = np.zeros((new_height, new_width, 3), dtype=np.uint8)
        # Place original image in the center
        result[:, pad_left:pad_left+w, :] = image_np
        
        # Create mask for inpainting
        inpaint_mask = np.ones((new_height, new_width), dtype=np.uint8) * 255
        inpaint_mask[:, pad_left:pad_left+w] = 0
        
        # Perform outpainting using Stable Diffusion
        result = outpaint_regions(result, inpaint_mask, prompt)
        
    else:
        # Need to add height (outpaint top/bottom)
        new_width = w
        new_height = int(w / target_ratio_value)
        
        # Calculate padding
        pad_height = new_height - h
        pad_top = pad_height // 2
        pad_bottom = pad_height - pad_top
        
        # Create canvas with padding
        result = np.zeros((new_height, new_width, 3), dtype=np.uint8)
        # Place original image in the center
        result[pad_top:pad_top+h, :, :] = image_np
        
        # Create mask for inpainting
        inpaint_mask = np.ones((new_height, new_width), dtype=np.uint8) * 255
        inpaint_mask[pad_top:pad_top+h, :] = 0
        
        # Perform outpainting using Stable Diffusion
        result = outpaint_regions(result, inpaint_mask, prompt)
    
    return result

def outpaint_regions(image, mask, prompt):
    """Use Stable Diffusion to outpaint masked regions"""
    # Convert to PIL images
    image_pil = Image.fromarray(image)
    mask_pil = Image.fromarray(mask)
    
    # If prompt is empty, use a generic one
    if not prompt or prompt.strip() == "":
        prompt = "seamless extension of the image, same style, same scene"
    
    # Generate the outpainting
    output = inpaint_model(
        prompt=prompt,
        image=image_pil,
        mask_image=mask_pil,
        guidance_scale=7.5,
        num_inference_steps=25
    ).images[0]
    
    return np.array(output)

def process_image(input_image, target_ratio="16:9", prompt=""):
    """Main processing function for the Gradio interface"""
    try:
        # Convert from Gradio format
        if isinstance(input_image, dict) and 'image' in input_image:
            image = input_image['image']
        else:
            image = input_image
            
        # Convert PIL to numpy if needed
        if isinstance(image, Image.Image):
            image_np = np.array(image)
        else:
            image_np = image
            
        # Get SAM mask to identify important regions
        mask = get_sam_mask(image_np)
        
        # Adjust aspect ratio while preserving content
        result = adjust_aspect_ratio(image_np, mask, target_ratio, prompt)
        
        # Convert result to PIL for visualization
        result_pil = Image.fromarray(result)
        
        return result_pil
    
    except Exception as e:
        print(f"Error processing image: {e}")
        return None

# Create the Gradio interface
with gr.Blocks(title="Automatic Aspect Ratio Adjuster") as demo:
    gr.Markdown("# Automatic Aspect Ratio Adjuster")
    gr.Markdown("Upload an image, choose your target aspect ratio, and let the AI adjust it while preserving important content.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image")
            
            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    choices=["16:9", "4:3", "1:1", "9:16", "3:4"], 
                    value="16:9",
                    label="Target Aspect Ratio"
                )
                
            prompt = gr.Textbox(
                label="Outpainting Prompt (optional)", 
                placeholder="Describe the scene for better outpainting"
            )
            
            submit_btn = gr.Button("Process Image")
        
        with gr.Column():
            output_image = gr.Image(label="Processed Image")
    
    submit_btn.click(
        process_image, 
        inputs=[input_image, aspect_ratio, prompt], 
        outputs=output_image
    )
    
    gr.Markdown("""
    ## How it works
    1. SAM (Segment Anything Model) identifies important content in your image
    2. The algorithm calculates how to adjust the aspect ratio while preserving this content
    3. Stable Diffusion fills in the new areas with AI-generated content that matches the original image
    
    ## Tips
    - For best results, provide a descriptive prompt that matches the scene
    - Try different aspect ratios to see what works best
    - The model works best with clear, well-lit images
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()