Spaces:
Sleeping
Sleeping
File size: 3,771 Bytes
ed8261c f171301 f0406ee f171301 f0406ee f171301 f0406ee f171301 f0406ee f171301 f0406ee f171301 ed8261c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import gradio as gr
import torch
from diffusers import StableDiffusionInpaintPipeline
import cv2
from PIL import Image
pipe = StableDiffusionInpaintPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-inpainting",
torch_dtype=torch.float16,
)
def zoom_out(image, n):
# Original dimensions
original_width, original_height = image.size
# Calculate new dimensions after shrinking
new_width = int(original_width // n)
new_height = int(original_height // n)
print(new_width,new_height)
# Resize the image (shrink the object)
resized_image = image.resize((new_width, new_height), Image.ANTIALIAS)
# Create a new white canvas that is n times larger than the original
new_canvas_width = original_width
new_canvas_height = original_height
new_image = Image.new("RGB", (new_canvas_width, new_canvas_height), "white")
# Calculate the position to paste the resized object in the center
x_offset = (new_canvas_width - new_width) // 2
y_offset = (new_canvas_height - new_height) // 2
# Paste the resized image onto the center of the new canvas
new_image.paste(resized_image, (x_offset, y_offset))
return new_image
def create_contour_mask(image, tolerance=10):
# Convert the image to a numpy array
image_array = np.array(image)
# Define a tolerance range around white
lower_bound = np.array([255 - tolerance] * 3)
upper_bound = np.array([255] * 3)
# Create a mask: black for background, white for object
mask = np.ones((image_array.shape[0], image_array.shape[1]), dtype=np.uint8) * 255
non_white_pixels = np.any((image_array < lower_bound) | (image_array > upper_bound), axis=-1)
mask[non_white_pixels] = 0 # Non-white pixels become black in the mask
# Perform dilation to fill in any gaps (like white text within the object)
kernel = np.ones((2, 2), np.uint8) # You can adjust the size of the kernel as needed
dilated_mask = cv2.dilate(mask, kernel, iterations=2)
# Invert the dilated mask to get the final object mask
# object_mask = 255 - dilated_mask
# Convert to PIL image for easier visualization and further processing
mask_image = Image.fromarray(dilated_mask)
return mask_image
def process_image(image: Image.Image, prompt: str, slider_value: float) -> Image.Image:
# Placeholder function for processing
# Replace this with your actual processing logic
# For example, modifying the image based on the slider value and prompt
if slider_value != 1:
image = zoom_out(image,slider_value)
mask = create_contour_mask(image,10)
# processed_image = image.copy() # Just returning a copy for now
processed_image = pipe(prompt=prompt, image=image_big, mask_image=contour_mask).images[0]
return processed_image
with gr.Blocks() as demo:
# Title at the top center
gr.Markdown("<h1 style='text-align: center;'>Imagine Backgrounds</h1>")
with gr.Row():
with gr.Column(scale=1):
# Image upload on the left
image_input = gr.Image(type='pil', label='Upload Image')
# Slider below the image upload
slider = gr.Slider(minimum=1, maximum=4, step=0.2, value=1, label='Select Zoom')
# Textbox for prompt
prompt_input = gr.Textbox(label='Enter Prompt')
# Submit button
submit_btn = gr.Button("Submit")
with gr.Column(scale=1):
# Output image on the right
image_output = gr.Image(label='Output Image')
# Event handler to process the image when the button is clicked
submit_btn.click(fn=process_image, inputs=[image_input, prompt_input, slider], outputs=image_output)
# Launch the Gradio app
demo.launch(debug=True)
|