fixed some bugs

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from transformers import AutoModel
 from transformers import SamModel, SamConfig, SamProcessor
 from PIL import Image
+import matplotlib.pyplot as plt
 
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -37,37 +38,42 @@ def get_bbox(gt_map):
     return bbox
 
 
-
-def greet(image):
-    image = Image.fromarray(image)
+def process_image(image_input):
+    # Convert the input to a PIL Image and resize
+    image = Image.fromarray(image_input).convert('RGB')
     image = image.resize((256, 256))
     
-    gt_mask = np.array(image)
-    prompt = get_bbox(gt_mask)
+    # Create a prompt based on the image size
+    prompt = [0, 0, image.width, image.height]
+    prompt = [[prompt]]  # Modify the prompt to be in the expected format for the processor
 
-    inputs = processor(images=image, input_boxes=[[prompt]], return_tensors="pt")
+    # Process the image and bounding box
+    inputs = processor(image, input_boxes=prompt, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
+    # Forward pass without gradient calculation
     model.eval()
     with torch.no_grad():
-        outputs = model(**inputs, multimask_outputs=False)
-    
-    seg_prob = torch.sigmoid(outputs.pred_masks.squeeze(0))
-    seg_prob = seg_prob.cpu().numpy().squeeze()
-    seg_prob = (seg_prob > 0.5).astype(np.uint8)
+        outputs = model(**inputs, multimask_output=False)
 
-    # Ensure the array is 2D (height, width) for grayscale image
-    if seg_prob.ndim > 2:
-        seg_prob = seg_prob.squeeze()  # Remove extra dimensions if any
-    elif seg_prob.ndim < 2:
-        raise ValueError("Output mask has less than 2 dimensions")
-
-    # Convert the processed mask back to a PIL image
-    seg_prob_image = Image.fromarray(seg_prob)
+    # Process model output
+    seg_prob = torch.sigmoid(outputs['pred_masks'].squeeze(1))
+    seg_prob = seg_prob.cpu().numpy().squeeze()
+    seg = (seg_prob > 0.5).astype(np.uint8)
 
-    return seg_prob_image
+    # Convert numpy arrays back to PIL Images for Gradio output
+    seg_image = Image.fromarray(seg * 255)  # Convert boolean mask to uint8 image
+    prob_map = Image.fromarray((seg_prob * 255).astype(np.uint8))  # Scale probabilities to 0-255
 
+    return image, seg_image, prob_map
 
+# Define Gradio interface
+iface = gr.Interface(
+    fn=process_image, 
+    inputs=gr.inputs.Image(shape=(256, 256)), 
+    outputs=[gr.outputs.Image(label="Original Image"), gr.outputs.Image(label="Segmentation Mask"), gr.outputs.Image(label="Probability Map")],
+    title="Image Segmentation"
+)
 
-iface = gr.Interface(fn= greet, inputs="image", outputs="image", title="Greeter")
-iface.launch()
+# Launch the interface
+iface.launch()