Update app.py

app.py

@@ -3,28 +3,16 @@ import gradio as gr
 from PIL import Image
 import torch
 import numpy as np
-import threading
 
+# Load CLIPSeg processor and model
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
-# Function to process image and generate mask
 # Function to process image and generate mask
 def process_image(image, prompt):
     inputs = processor(
         text=prompt, images=image, padding="max_length", return_tensors="pt"
     )
-
-    # Extract image tensor and normalize it
-    image_tensor = inputs["pixel_values"].squeeze().permute(1, 2, 0).cpu().numpy()
-    image_tensor = (image_tensor * 255).astype(np.uint8)
-    image_tensor = Image.fromarray(image_tensor)
-    image_tensor = image_tensor.convert("RGB")
-    
-    # Perform CLIPSeg processing
-    inputs = processor(
-        text=prompt, images=image_tensor, padding="max_length", return_tensors="pt"
-    )
     with torch.no_grad():
         outputs = model(**inputs)
         preds = outputs.logits
@@ -42,6 +30,16 @@ def process_image(image, prompt):
 
     return mask
 
+# Function to get masks from positive or negative prompts
+def get_masks(prompts, img, threshold):
+    prompts = prompts.split(",")
+    masks = []
+    for prompt in prompts:
+        mask = process_image(img, prompt)
+        mask = mask > threshold
+        masks.append(mask)
+
+    return masks
 
 # Function to extract image using positive and negative prompts
 def extract_image(pos_prompts, neg_prompts, img, threshold):
@@ -58,6 +56,7 @@ def extract_image(pos_prompts, neg_prompts, img, threshold):
 
     return output_image, final_mask
 
+# Define Gradio interface
 # Define Gradio interface
 iface = gr.Interface(
     fn=extract_image,
@@ -70,14 +69,18 @@ iface = gr.Interface(
             label="Please describe what you want to ignore (comma separated)",
             key="neg_prompts",
         ),
-        gr.Image(type="pil", label="Input Image", key="img"),
-        gr.Slider(minimum=0, maximum=1, default=0.4, label="Threshold", key="threshold"),
+        gr.Image(type="pil", label="Input Image"),
+        gr.Slider(minimum=0, maximum=1, default=0.4, label="Threshold"),
     ],
     outputs=[
-        gr.Image(label="Result", key="output_image"),
-        gr.Image(label="Mask", key="output_mask"),
+        gr.Image(label="Result"),
+        gr.Image(label="Mask"),
     ],
 )
 
 # Launch Gradio API
 iface.launch()
+
+
+# Launch Gradio API
+iface.launch()