try np fix when using other models from default
Browse files
app.py
CHANGED
@@ -1,6 +1,7 @@
|
|
1 |
import gradio as gr
|
2 |
from rembg import remove, new_session
|
3 |
from PIL import Image
|
|
|
4 |
import logging
|
5 |
import time
|
6 |
|
@@ -14,7 +15,7 @@ MODEL_OPTIONS = {
|
|
14 |
"isnet-general-use": "A new pre-trained model for general use cases",
|
15 |
"isnet-anime": "High-accuracy segmentation for anime characters",
|
16 |
"silueta": "A reduced-size version of u2net (43MB)",
|
17 |
-
"sam": "A pre-trained model for any use case",
|
18 |
"unet": "Lightweight version of u2net model",
|
19 |
"u2netp": "A lightweight version of u2net model",
|
20 |
"u2net_human_seg": "A pre-trained model for human segmentation",
|
@@ -58,14 +59,21 @@ def remove_background(input_image, bg_color, model_choice, alpha_matting, post_p
|
|
58 |
|
59 |
logging.info(f'Model name={model_name}')
|
60 |
logging.info(remove_kwargs)
|
|
|
|
|
|
|
|
|
61 |
# Use the remove function
|
62 |
if session or bg_color_rgba:
|
63 |
-
|
64 |
else:
|
65 |
-
|
66 |
|
67 |
logging.info("Background removed")
|
68 |
|
|
|
|
|
|
|
69 |
# Convert to RGB mode if necessary
|
70 |
if not only_mask and output_image.mode != 'RGB':
|
71 |
output_image = output_image.convert('RGB')
|
@@ -74,7 +82,7 @@ def remove_background(input_image, bg_color, model_choice, alpha_matting, post_p
|
|
74 |
# Save the output image to a temporary file
|
75 |
# Generate a unique timestamp for the output file name
|
76 |
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
77 |
-
output_path = f"
|
78 |
output_image.save(output_path)
|
79 |
logging.info(f"Saved output image {output_path}")
|
80 |
|
@@ -104,7 +112,7 @@ iface = gr.Interface(
|
|
104 |
gr.Image(type="pil", label="Output Image"),
|
105 |
gr.File(label="Download the output image")
|
106 |
],
|
107 |
-
title="Advanced Background Remover v2.
|
108 |
description="Upload an image to remove the background. Customize the result with different options, including background color, model selection, alpha matting, and more.",
|
109 |
allow_flagging="never",
|
110 |
)
|
|
|
1 |
import gradio as gr
|
2 |
from rembg import remove, new_session
|
3 |
from PIL import Image
|
4 |
+
import numpy as np
|
5 |
import logging
|
6 |
import time
|
7 |
|
|
|
15 |
"isnet-general-use": "A new pre-trained model for general use cases",
|
16 |
"isnet-anime": "High-accuracy segmentation for anime characters",
|
17 |
"silueta": "A reduced-size version of u2net (43MB)",
|
18 |
+
"sam": "A pre-trained model for any use case (Segment Anything Model)",
|
19 |
"unet": "Lightweight version of u2net model",
|
20 |
"u2netp": "A lightweight version of u2net model",
|
21 |
"u2net_human_seg": "A pre-trained model for human segmentation",
|
|
|
59 |
|
60 |
logging.info(f'Model name={model_name}')
|
61 |
logging.info(remove_kwargs)
|
62 |
+
|
63 |
+
# Convert PIL Image to numpy array
|
64 |
+
input_array = np.array(input_image)
|
65 |
+
|
66 |
# Use the remove function
|
67 |
if session or bg_color_rgba:
|
68 |
+
output_array = remove(input_array, **{k: v for k, v in remove_kwargs.items() if v is not None})
|
69 |
else:
|
70 |
+
output_array = remove(input_array) # Use the default remove function
|
71 |
|
72 |
logging.info("Background removed")
|
73 |
|
74 |
+
# Convert numpy array back to PIL Image
|
75 |
+
output_image = Image.fromarray(output_array)
|
76 |
+
|
77 |
# Convert to RGB mode if necessary
|
78 |
if not only_mask and output_image.mode != 'RGB':
|
79 |
output_image = output_image.convert('RGB')
|
|
|
82 |
# Save the output image to a temporary file
|
83 |
# Generate a unique timestamp for the output file name
|
84 |
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
85 |
+
output_path = f"output_remove_background_s{timestamp}.png"
|
86 |
output_image.save(output_path)
|
87 |
logging.info(f"Saved output image {output_path}")
|
88 |
|
|
|
112 |
gr.Image(type="pil", label="Output Image"),
|
113 |
gr.File(label="Download the output image")
|
114 |
],
|
115 |
+
title="Advanced Background Remover v2.6",
|
116 |
description="Upload an image to remove the background. Customize the result with different options, including background color, model selection, alpha matting, and more.",
|
117 |
allow_flagging="never",
|
118 |
)
|