Create app.py

app.py

@@ -0,0 +1,216 @@
+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()