codetechdevx commited on
Commit
d58ef1f
·
1 Parent(s): 9352bfd

added initial code

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py CHANGED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from loadimg import load_img
3
+ import spaces
4
+ from transformers import AutoModelForImageSegmentation
5
+ import torch
6
+ from torchvision import transforms
7
+
8
+ torch.set_float32_matmul_precision(["high", "highest"][0])
9
+
10
+ birefnet = AutoModelForImageSegmentation.from_pretrained(
11
+ "ZhengPeng7/BiRefNet", trust_remote_code=True
12
+ )
13
+ birefnet.to("cuda")
14
+
15
+ transform_image = transforms.Compose(
16
+ [
17
+ transforms.Resize((1024, 1024)),
18
+ transforms.ToTensor(),
19
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
20
+ ]
21
+ )
22
+
23
+ def fn(image):
24
+ im = load_img(image, output_type="pil")
25
+ im = im.convert("RGB")
26
+ origin = im.copy()
27
+ processed_image = process(im)
28
+ return (processed_image, origin)
29
+
30
+ @spaces.GPU
31
+ def process(image):
32
+ image_size = image.size
33
+ input_images = transform_image(image).unsqueeze(0).to("cuda")
34
+ # Prediction
35
+ with torch.no_grad():
36
+ preds = birefnet(input_images)[-1].sigmoid().cpu()
37
+ pred = preds[0].squeeze()
38
+ pred_pil = transforms.ToPILImage()(pred)
39
+ mask = pred_pil.resize(image_size)
40
+ image.putalpha(mask)
41
+ return image
42
+
43
+ def process_file(f):
44
+ name_path = f.rsplit(".", 1)[0] + ".png"
45
+ im = load_img(f, output_type="pil")
46
+ im = im.convert("RGB")
47
+ transparent = process(im)
48
+ transparent.save(name_path)
49
+ return name_path
50
+
51
+ slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
52
+ slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
53
+ image_upload = gr.Image(label="Upload an image")
54
+ image_file_upload = gr.Image(label="Upload an image", type="filepath")
55
+ url_input = gr.Textbox(label="Paste an image URL")
56
+ output_file = gr.File(label="Output PNG File")
57
+
58
+ # Example images
59
+ chameleon = load_img("butterfly.jpg", output_type="pil")
60
+ url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
61
+
62
+ tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
63
+ tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
64
+ tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
65
+
66
+ demo = gr.TabbedInterface(
67
+ [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
68
+ )
69
+
70
+ if __name__ == "__main__":
71
+ demo.launch(show_error=True)