Update app.py

app.py

@@ -8,11 +8,23 @@ import threading
 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
@@ -30,16 +42,6 @@ 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):