File size: 2,441 Bytes
a55dfd9
 
63abc5c
fd49265
8f3a72a
fd49265
8f3a72a
7910b4e
25a04e5
8f3a72a
 
a55dfd9
8f3a72a
7910b4e
8f3a72a
 
d76bdf0
7910b4e
fd49265
8f3a72a
d76bdf0
 
fd49265
8f3a72a
 
 
 
fd49265
7910b4e
 
a55dfd9
8f3a72a
5a72035
94e9c31
8f3a72a
5a72035
fd49265
8f3a72a
 
7910b4e
8f3a72a
fd49265
 
8f3a72a
fd49265
8f3a72a
 
 
fd49265
 
8f3a72a
 
 
 
 
 
 
fd49265
94e9c31
8f3a72a
fd49265
94e9c31
8f3a72a
fd49265
8f3a72a
 
 
 
a55dfd9
19dc428
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
import gradio as gr
from rembg import remove
from PIL import Image, ImageDraw
import numpy as np
import io

# Function to remove background automatically
def remove_bg(image):
    if image is None:
        return None
    image = Image.open(image)
    output = remove(image)
    return output

# Function to refine background using + and - points
def refine_bg(image, points, threshold, mode):
    if image is None:
        return None
    
    image = Image.open(image).convert("RGBA")
    mask = Image.new("L", image.size, 0)
    draw = ImageDraw.Draw(mask)
    
    for point in points:
        x, y = point
        color = 255 if mode == "keep" else 0
        draw.ellipse((x - threshold * 10, y - threshold * 10, x + threshold * 10, y + threshold * 10), fill=color)
    
    refined = Image.composite(image, Image.new("RGBA", image.size, (0, 0, 0, 0)), mask)
    return refined

# Gradio UI
with gr.Blocks() as iface:
    gr.Markdown("# AI Background Remover")
    
    with gr.Row():
        input_img = gr.Image(type="filepath", label="Upload Image")
        output_img = gr.Image(type="pil", label="Output Image")
    
    remove_btn = gr.Button("Remove Background")
    remove_btn.click(remove_bg, inputs=input_img, outputs=output_img)
    
    refine_btn = gr.Button("Refine")
    refine_options = gr.Column(visible=False)
    
    with refine_options:
        gr.Markdown("### Refine Background")
        refine_editor = gr.Image(type="filepath", label="Tap to Add Points", interactive=True)
        keep_btn = gr.Button("+ (Keep)")
        remove_btn = gr.Button("- (Remove)")
        threshold_slider = gr.Slider(0.00, 1.00, value=0.5, step=0.01, label="Threshold")
        apply_refine_btn = gr.Button("Apply Refinements")
    
    refine_mode = gr.State("keep")
    
    def show_refine_section(image):
        return gr.update(visible=True), image  # Show section and keep image visible
    
    def set_mode_keep():
        return "keep", "+ (Keep) Selected"
    
    def set_mode_remove():
        return "remove", "- (Remove) Selected"
    
    refine_btn.click(show_refine_section, inputs=input_img, outputs=[refine_options, refine_editor])
    keep_btn.click(set_mode_keep, outputs=[refine_mode, keep_btn])
    remove_btn.click(set_mode_remove, outputs=[refine_mode, remove_btn])
    apply_refine_btn.click(refine_bg, inputs=[input_img, refine_editor, threshold_slider, refine_mode], outputs=output_img)

iface.launch()