Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -7,13 +7,13 @@ from torchvision import transforms
|
|
| 7 |
import uuid
|
| 8 |
import os
|
| 9 |
|
| 10 |
-
#
|
| 11 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 12 |
print(f"Using device: {device}")
|
| 13 |
|
| 14 |
torch.set_float32_matmul_precision(["high", "highest"][0])
|
| 15 |
|
| 16 |
-
# Load model
|
| 17 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
| 18 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
| 19 |
)
|
|
@@ -40,60 +40,24 @@ def process(image):
|
|
| 40 |
image.putalpha(mask)
|
| 41 |
return image
|
| 42 |
|
| 43 |
-
#
|
| 44 |
def fn(image):
|
| 45 |
im = load_img(image, output_type="pil").convert("RGB")
|
| 46 |
processed_image = process(im)
|
|
|
|
| 47 |
filename = f"/tmp/processed_{uuid.uuid4().hex}.png"
|
| 48 |
processed_image.save(filename)
|
| 49 |
-
return processed_image, filename
|
| 50 |
-
|
| 51 |
-
# File tab
|
| 52 |
-
def process_file(f):
|
| 53 |
-
name_path = f.rsplit(".", 1)[0] + ".png"
|
| 54 |
-
im = load_img(f, output_type="pil").convert("RGB")
|
| 55 |
-
transparent = process(im)
|
| 56 |
-
transparent.save(name_path)
|
| 57 |
-
return name_path
|
| 58 |
|
| 59 |
-
|
| 60 |
-
url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
|
| 61 |
-
chameleon = load_img(url_example, output_type="pil")
|
| 62 |
|
| 63 |
-
#
|
| 64 |
-
|
| 65 |
fn,
|
| 66 |
inputs=gr.Image(label="Upload an image"),
|
| 67 |
outputs=[
|
| 68 |
-
gr.Image(label="Preview"),
|
| 69 |
-
gr.File(label="Download
|
| 70 |
],
|
| 71 |
-
examples=[chameleon],
|
| 72 |
-
api_name="image"
|
| 73 |
-
)
|
| 74 |
-
|
| 75 |
-
# Tab 2: URL input + Preview
|
| 76 |
-
tab2 = gr.Interface(
|
| 77 |
-
fn,
|
| 78 |
-
inputs=gr.Textbox(label="Paste an image URL"),
|
| 79 |
-
outputs=gr.Image(label="Preview"),
|
| 80 |
-
examples=[url_example],
|
| 81 |
-
api_name="text"
|
| 82 |
-
)
|
| 83 |
-
|
| 84 |
-
# Tab 3: File path input + downloadable result
|
| 85 |
-
tab3 = gr.Interface(
|
| 86 |
-
process_file,
|
| 87 |
-
inputs=gr.Image(label="Upload an image", type="filepath"),
|
| 88 |
-
outputs=gr.File(label="Output PNG File"),
|
| 89 |
-
examples=[],
|
| 90 |
-
api_name="png"
|
| 91 |
-
)
|
| 92 |
-
|
| 93 |
-
# Final App
|
| 94 |
-
demo = gr.TabbedInterface(
|
| 95 |
-
[tab1, tab2, tab3],
|
| 96 |
-
["Image Upload", "URL Input", "File Output"],
|
| 97 |
title="Background Removal Tool"
|
| 98 |
)
|
| 99 |
|
|
|
|
| 7 |
import uuid
|
| 8 |
import os
|
| 9 |
|
| 10 |
+
# Select device
|
| 11 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 12 |
print(f"Using device: {device}")
|
| 13 |
|
| 14 |
torch.set_float32_matmul_precision(["high", "highest"][0])
|
| 15 |
|
| 16 |
+
# Load BiRefNet model
|
| 17 |
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
| 18 |
"ZhengPeng7/BiRefNet", trust_remote_code=True
|
| 19 |
)
|
|
|
|
| 40 |
image.putalpha(mask)
|
| 41 |
return image
|
| 42 |
|
| 43 |
+
# Main function: image upload → preview + downloadable PNG
|
| 44 |
def fn(image):
|
| 45 |
im = load_img(image, output_type="pil").convert("RGB")
|
| 46 |
processed_image = process(im)
|
| 47 |
+
|
| 48 |
filename = f"/tmp/processed_{uuid.uuid4().hex}.png"
|
| 49 |
processed_image.save(filename)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 50 |
|
| 51 |
+
return processed_image, filename
|
|
|
|
|
|
|
| 52 |
|
| 53 |
+
# Gradio interface
|
| 54 |
+
demo = gr.Interface(
|
| 55 |
fn,
|
| 56 |
inputs=gr.Image(label="Upload an image"),
|
| 57 |
outputs=[
|
| 58 |
+
gr.Image(label="Processed Preview"),
|
| 59 |
+
gr.File(label="Download PNG")
|
| 60 |
],
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 61 |
title="Background Removal Tool"
|
| 62 |
)
|
| 63 |
|