Spaces:
Running
on
L40S
Running
on
L40S
Update app.py
Browse files
app.py
CHANGED
@@ -36,22 +36,23 @@ images_examples = [
|
|
36 |
if os.path.isfile(os.path.join(examples_folder, file))
|
37 |
]
|
38 |
|
39 |
-
def remove_background(
|
40 |
-
|
41 |
# Create a temporary folder for downloaded and processed images
|
42 |
temp_dir = tempfile.mkdtemp()
|
43 |
-
|
44 |
-
image_path = os.path.join(temp_dir, 'input_image.png')
|
45 |
unique_id = str(uuid.uuid4())
|
46 |
-
|
47 |
-
|
48 |
try:
|
49 |
# Check if input_url is already a PIL Image
|
50 |
-
if isinstance(
|
51 |
-
image =
|
52 |
else:
|
53 |
# Otherwise, assume it's a file path and open it
|
54 |
-
image = Image.open(
|
|
|
|
|
|
|
55 |
|
56 |
# Save the resized image
|
57 |
image.save(image_path)
|
@@ -61,6 +62,7 @@ def remove_background(input_url):
|
|
61 |
|
62 |
if remove_bg is True:
|
63 |
# Run background removal
|
|
|
64 |
try:
|
65 |
img = Image.open(image_path)
|
66 |
result = remove(img)
|
@@ -110,12 +112,12 @@ def run_inference(temp_dir, removed_bg_path):
|
|
110 |
except subprocess.CalledProcessError as e:
|
111 |
return f"Error during inference: {str(e)}"
|
112 |
|
113 |
-
def process_image(
|
114 |
|
115 |
torch.cuda.empty_cache()
|
116 |
|
117 |
# Remove background
|
118 |
-
result = remove_background(
|
119 |
|
120 |
if isinstance(result, str) and result.startswith("Error"):
|
121 |
raise gr.Error(f"{result}") # Return the error message if something went wrong
|
@@ -166,8 +168,11 @@ def gradio_interface():
|
|
166 |
image_mode="RGBA",
|
167 |
height=240
|
168 |
)
|
|
|
|
|
169 |
|
170 |
submit_button = gr.Button("Process")
|
|
|
171 |
gr.Examples(
|
172 |
examples = examples_folder,
|
173 |
inputs = [input_image],
|
@@ -176,7 +181,7 @@ def gradio_interface():
|
|
176 |
|
177 |
output_video= gr.Video(label="Output Video", scale=4)
|
178 |
|
179 |
-
submit_button.click(process_image, inputs=[input_image], outputs=[output_video])
|
180 |
|
181 |
return app
|
182 |
|
|
|
36 |
if os.path.isfile(os.path.join(examples_folder, file))
|
37 |
]
|
38 |
|
39 |
+
def remove_background(input_pil, remove_bg):
|
40 |
+
|
41 |
# Create a temporary folder for downloaded and processed images
|
42 |
temp_dir = tempfile.mkdtemp()
|
|
|
|
|
43 |
unique_id = str(uuid.uuid4())
|
44 |
+
image_path = os.path.join(temp_dir, f'input_image_{unique_id}.png')
|
45 |
+
|
46 |
try:
|
47 |
# Check if input_url is already a PIL Image
|
48 |
+
if isinstance(input_pil, Image.Image):
|
49 |
+
image = input_pil
|
50 |
else:
|
51 |
# Otherwise, assume it's a file path and open it
|
52 |
+
image = Image.open(input_pil)
|
53 |
+
|
54 |
+
# Flip the image horizontally
|
55 |
+
image = image.transpose(Image.FLIP_LEFT_RIGHT)
|
56 |
|
57 |
# Save the resized image
|
58 |
image.save(image_path)
|
|
|
62 |
|
63 |
if remove_bg is True:
|
64 |
# Run background removal
|
65 |
+
removed_bg_path = os.path.join(temp_dir, f'output_image_rmbg_{unique_id}.png')
|
66 |
try:
|
67 |
img = Image.open(image_path)
|
68 |
result = remove(img)
|
|
|
112 |
except subprocess.CalledProcessError as e:
|
113 |
return f"Error during inference: {str(e)}"
|
114 |
|
115 |
+
def process_image(input_pil, remove_bg):
|
116 |
|
117 |
torch.cuda.empty_cache()
|
118 |
|
119 |
# Remove background
|
120 |
+
result = remove_background(input_pil, remove_bg)
|
121 |
|
122 |
if isinstance(result, str) and result.startswith("Error"):
|
123 |
raise gr.Error(f"{result}") # Return the error message if something went wrong
|
|
|
168 |
image_mode="RGBA",
|
169 |
height=240
|
170 |
)
|
171 |
+
|
172 |
+
remove_bg = gr.Checkbox(label="Need to remove BG ?", value=False)
|
173 |
|
174 |
submit_button = gr.Button("Process")
|
175 |
+
|
176 |
gr.Examples(
|
177 |
examples = examples_folder,
|
178 |
inputs = [input_image],
|
|
|
181 |
|
182 |
output_video= gr.Video(label="Output Video", scale=4)
|
183 |
|
184 |
+
submit_button.click(process_image, inputs=[input_image, remove_bg], outputs=[output_video])
|
185 |
|
186 |
return app
|
187 |
|