Spaces:
Running
Running
File size: 3,531 Bytes
5a3dfd3 6dcded2 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a 69f51b6 24e5a6a |
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 100 |
import os
import imageio
from PIL import Image
import gradio as gr
import cv2
import paddlehub as hub
import onnxruntime
# Download and setup models
os.system("wget https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx")
os.system("pip install onnxruntime imageio")
os.makedirs("data", exist_ok=True)
os.makedirs("dataout", exist_ok=True)
# Load LaMa ONNX model
sess_options = onnxruntime.SessionOptions()
lama_model = onnxruntime.InferenceSession('lama_fp32.onnx', sess_options=sess_options)
# Load U^2-Net model for automatic masking
u2net_model = hub.Module(name='U2Net')
# --- Helper Functions ---
def prepare_image(image, target_size=(512, 512)):
"""Resizes and preprocesses image for LaMa model."""
if isinstance(image, Image.Image):
image = image.resize(target_size)
image = np.array(image)
elif isinstance(image, np.ndarray):
image = cv2.resize(image, target_size)
else:
raise ValueError("Input image should be either PIL Image or numpy array!")
# Normalize to [0, 1] and convert to CHW format
image = image.astype(np.float32) / 255.0
if image.ndim == 3:
image = np.transpose(image, (2, 0, 1))
elif image.ndim == 2:
image = image[np.newaxis, ...]
return image[np.newaxis, ...] # Add batch dimension
def generate_mask(image, method="automatic"):
"""Generates mask from image using U^2-Net or user input."""
if method == "automatic":
input_size = 320 # Adjust based on U^2-Net requirements
result = u2net_model.Segmentation(
images=[cv2.cvtColor(image, cv2.COLOR_RGB2BGR)],
paths=None,
batch_size=1,
input_size=input_size,
output_dir='output',
visualization=False
)
mask = Image.fromarray(result[0]['mask'])
mask = mask.resize((512, 512)) # Resize to match LaMa input
mask.save("./data/data_mask.png")
else: # "manual"
mask = imageio.imread("./data/data_mask.png")
mask = Image.fromarray(mask).convert("L") # Ensure grayscale
mask = mask.resize((512, 512))
return prepare_image(mask, (512, 512))
def inpaint_image(image, mask):
"""Performs inpainting using the LaMa model."""
outputs = lama_model.run(None, {'image': image, 'mask': mask})
output = outputs[0][0]
output = output.transpose(1, 2, 0)
output = (output * 255).astype(np.uint8)
return Image.fromarray(output)
# --- Gradio Interface ---
def process_image(input_image, mask_option):
"""Main function for Gradio interface."""
imageio.imwrite("./data/data.png", input_image)
image = prepare_image(input_image)
mask = generate_mask(input_image, method=mask_option)
inpainted_image = inpaint_image(image, mask)
inpainted_image = inpainted_image.resize(Image.open("./data/data.png").size)
inpainted_image.save("./dataout/data_mask.png")
return "./dataout/data_mask.png", "./data/data_mask.png"
iface = gr.Interface(
fn=process_image,
inputs=[
gr.Image(label="Input Image", type="numpy"),
gr.Radio(choices=["automatic", "manual"],
type="value", default="manual", label="Masking Option")
],
outputs=[
gr.Image(type="file", label="Inpainted Image"),
gr.Image(type="file", label="Generated Mask")
],
title="LaMa Image Inpainting",
description="Image inpainting with LaMa and U^2-Net. Upload your image and choose automatic or manual masking.",
)
iface.launch() |