rdjarbeng commited on
Commit
8423e33
·
1 Parent(s): 401fd9e

try np fix when using other models from default

Browse files
Files changed (1) hide show
  1. app.py +13 -5
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
- output_image = remove(input_image, **{k: v for k, v in remove_kwargs.items() if v is not None})
64
  else:
65
- output_image = remove(input_image) # Use the default remove function
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"output_remove_background_{timestamp}.png"
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.5",
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
  )