Spaces:
Sleeping
Sleeping
| 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() |