gaur3009/GraphicsAI

1import gradio as gr2from gradio_imageslider import ImageSlider3from PIL import Image, ImageDraw, ImageFont4import numpy as np5import cv26import torch7from torchvision import transforms8from transformers import AutoModelForImageSegmentation9 10torch.set_float32_matmul_precision(["high", "highest"][0])11 12# Load BiRefNet model for background removal13birefnet = AutoModelForImageSegmentation.from_pretrained(14    "ZhengPeng7/BiRefNet", trust_remote_code=True15)16birefnet.to("cuda")17transform_image = transforms.Compose(18    [19        transforms.Resize((1024, 1024)),20        transforms.ToTensor(),21        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),22    ]23)24 25def load_img(image, output_type="numpy"):26    if output_type == "pil":27        return Image.open(image).convert("RGB")28    else:29        return np.array(Image.open(image).convert("RGB"))30 31def add_text_to_image(image, text, position, color, font_size):32    img = Image.fromarray(image)33    draw = ImageDraw.Draw(img)34    font = ImageFont.truetype("arial.ttf", font_size)35    draw.text(position, text, fill=color, font=font)36    return np.array(img)37 38def inpaint_image(image, mask, inpaint_radius):39    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)40    mask_cv = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)41    result = cv2.inpaint(img_cv, mask_cv, inpaint_radius, cv2.INPAINT_TELEA)42    return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)43 44def background_removal(image):45    im = load_img(image, output_type="pil")46    im = im.convert("RGB")47    image_size = im.size48    origin = im.copy()49    image = load_img(im)50    input_images = transform_image(image).unsqueeze(0).to("cuda")51    with torch.no_grad():52        preds = birefnet(input_images)[-1].sigmoid().cpu()53    pred = preds[0].squeeze()54    pred_pil = transforms.ToPILImage()(pred)55    mask = pred_pil.resize(image_size)56    im.putalpha(mask)57    return (im, origin)58 59def update_image(image, text, color, font_size, mask_image, inpaint_radius):60    img_with_text = add_text_to_image(image, text, (50, 50), color, font_size)61    if mask_image is not None:62        mask = np.array(mask_image)63        img_with_text = inpaint_image(img_with_text, mask, inpaint_radius)64    return img_with_text65 66def fn(image):67    return background_removal(image)68 69slider1 = ImageSlider(label="Original Image", type="pil")70slider2 = ImageSlider(label="Processed Image", type="pil")71 72image_input = gr.Image(label="Upload an image for background removal")73text_input = gr.Textbox(label="Enter Text to Add", placeholder="Your text here...")74color_input = gr.ColorPicker(label="Text Color")75font_size_input = gr.Slider(minimum=10, maximum=100, label="Font Size")76mask_input = gr.Image(type="numpy", label="Upload Mask Image (for Inpainting)", optional=True)77inpaint_radius_input = gr.Slider(minimum=1, maximum=50, value=3, label="Inpaint Radius")78 79bg_removal_interface = gr.Interface(80    fn, inputs=image_input, outputs=slider1, examples=["chameleon.jpg"]81)82 83design_editing_interface = gr.Interface(84    fn=lambda image, text, color, font_size, mask_image, inpaint_radius: update_image(image, text, color, font_size, mask_image, inpaint_radius),85    inputs=[image_input, text_input, color_input, font_size_input, mask_input, inpaint_radius_input],86    outputs=slider287)88 89demo = gr.TabbedInterface(90    [bg_removal_interface, design_editing_interface],91    ["Background Removal", "Design Editing"],92    title="Advanced Image Editor"93)94 95if __name__ == "__main__":96    demo.launch()97