Update app.py

app.py

@@ -16,9 +16,8 @@ import gradio as gr
 from huggingface_hub import hf_hub_download, snapshot_download
 from ip_adapter import IPAdapterXL
 from safetensors.torch import load_file
-# remove the background
-from tsr.system import TSR
-from tsr.utils import remove_background, resize_foreground
+from torchvision.models.detection import deeplabv3_resnet101
+from torchvision.transforms import functional as F
 
 snapshot_download(
     repo_id="h94/IP-Adapter", allow_patterns="sdxl_models/*", local_dir="."
@@ -103,31 +102,8 @@ def resize_img(
             np.array(input_image)
         )
         input_image = Image.fromarray(res)
-    # return input_image
-    return resized_image
-
-    #added
-
-def preprocess(resized_image, do_remove_background, foreground_ratio):
-    def fill_background(image):
-        image = np.array(image).astype(np.float32) / 255.0
-        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
-        image = Image.fromarray((image * 255.0).astype(np.uint8))
-        return image
-
-    if do_remove_background:
-        image = input_image.convert("RGB")
-        image = remove_background(image, rembg_session)
-        image = resize_foreground(image, foreground_ratio)
-        image = fill_background(image)
-    else:
-        image = input_image
-        if image.mode == "RGBA":
-            image = fill_background(image)
     return input_image
 
-    #/added
-
 examples = [
     [
         "./asset/0.jpg",
@@ -162,6 +138,43 @@ def run_for_examples(style_image, source_image, prompt, scale, control_scale):
         neg_content_scale=0,
     )
 
+# Add the background removal function
+def remove_background(input_image):
+    # Load the deep learning model
+    model = deeplabv3_resnet101(pretrained=True)
+    model.eval()
+
+    # Preprocess the image
+    preprocess = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    input_tensor = preprocess(input_image)
+    input_batch = input_tensor.unsqueeze(0)  # Create a mini-batch as expected by the model
+
+    # Move the input and model to GPU for speed if available
+    if torch.cuda.is_available():
+        input_batch = input_batch.to('cuda')
+        model.to('cuda')
+
+    with torch.no_grad():
+        output = model(input_batch)['out'][0]
+    output_predictions = output.argmax(0)
+
+    # Create a binary (black and white) mask of the profile foreground
+    mask = output_predictions.byte().cpu().numpy()
+    background = np.zeros(mask.shape)
+    bin_mask = np.where(mask, 255, background).astype(np.uint8)
+
+    # Create a transparent foreground
+    b, g, r = cv2.split(np.array(input_image).astype('uint8'))
+    a = np.ones(bin_mask.shape, dtype='uint8') * 255
+    alpha_im = cv2.merge([b, g, r, a], 4)
+    bg = np.zeros(alpha_im.shape)
+    new_mask = np.stack([bin_mask, bin_mask, bin_mask, bin_mask], axis=2)
+    foreground = np.where(new_mask, alpha_im, bg).astype(np.uint8)
+
+    return foreground
 
 @spaces.GPU
 def create_image(
@@ -208,6 +221,17 @@ def create_image(
         cv_input_image = pil_to_cv2(input_image)
         detected_map = cv2.Canny(cv_input_image, 50, 200)
         canny_map = Image.fromarray(cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB))
+        
+        # Remove background from the input image
+        foreground = remove_background(input_image)
+        # Convert the foreground back to a PIL image if necessary
+        foreground_pil = Image.fromarray(foreground)
+        
+        # Use foreground_pil instead of input_image for further processing
+        # Note: You might need to adjust the following lines based on how you intend to use the foreground_pil
+        # For example, if you're passing it to the IP-Adapter, ensure it's in the correct format
+        
+        # Continue with the existing logic for generating the image...
     else:
         canny_map = Image.new("RGB", (1024, 1024), color=(255, 255, 255))
         control_scale = 0
@@ -215,7 +239,7 @@ def create_image(
     if float(control_scale) == 0:
         canny_map = canny_map.resize((1024, 1024))
 
-    if len(neg_content_prompt) > 0 and neg_content_scale != 0:
+    if len(neg_content_prompt) > 0 and neg_content_scale!= 0:
         images = ip_model.generate(
             pil_image=image_pil,
             prompt=prompt,
@@ -282,7 +306,7 @@ with block:
                 with gr.Row():
                     with gr.Column():
                         image_pil = gr.Image(label="Style Image", type="pil")
-                        processed_image = gr.Image(label="Preprocess uWu", interactive=False)
+                        # processed_image = gr.Image(label="Preprocess uWu", interactive=False)
                     with gr.Column():
                         prompt = gr.Textbox(
                             label="Prompt",