fffiloni commited on
Commit
d63b039
1 Parent(s): 39c37ca

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -12
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(input_url):
40
- remove_bg = True
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
- removed_bg_path = os.path.join(temp_dir, f'output_image_rmbg_{unique_id}.png')
47
-
48
  try:
49
  # Check if input_url is already a PIL Image
50
- if isinstance(input_url, Image.Image):
51
- image = input_url
52
  else:
53
  # Otherwise, assume it's a file path and open it
54
- image = Image.open(input_url)
 
 
 
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(input_url):
114
 
115
  torch.cuda.empty_cache()
116
 
117
  # Remove background
118
- result = remove_background(input_url)
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