Update app.py

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, jsonify, render_template
+from flask import Flask, request, jsonify
 from PIL import Image
 import base64
 from io import BytesIO
@@ -13,7 +13,7 @@ app = Flask(__name__)
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
-def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
+def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
     inputs = processor(
         text=prompt, images=image, padding="max_length", return_tensors="pt"
     )
@@ -26,7 +26,6 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
     pred = torch.sigmoid(preds)
     mat = pred.cpu().numpy()
     mask = Image.fromarray(np.uint8(mat * 255), "L")
-    mask = mask.convert("RGB")
     mask = mask.resize(image.size)
     mask = np.array(mask)[:, :, 0]
 
@@ -36,34 +35,14 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
     mask = (mask - mask_min) / (mask_max - mask_min)
 
     # threshold the mask
-    bmask = mask > threhsold
+    bmask = mask > threshold
     # zero out values below the threshold
-    mask[mask < threhsold] = 0
-
-    fig, ax = plt.subplots()
-    ax.imshow(image)
-    ax.imshow(mask, alpha=alpha_value, cmap="jet")
-
-    if draw_rectangles:
-        contours, hierarchy = cv2.findContours(
-            bmask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
-        )
-        for contour in contours:
-            x, y, w, h = cv2.boundingRect(contour)
-            rect = plt.Rectangle(
-                (x, y), w, h, fill=False, edgecolor="yellow", linewidth=2
-            )
-            ax.add_patch(rect)
-
-    ax.axis("off")
-    plt.tight_layout()
+    mask[mask < threshold] = 0
 
     bmask = Image.fromarray(bmask.astype(np.uint8) * 255, "L")
-    output_image = Image.new("RGBA", image.size, (0, 0, 0, 0))
-    output_image.paste(image, mask=bmask)
 
-    return fig, mask, output_image
-    
+    return bmask
+
 @app.route('/')
 def index():
     return "Hello, World! clipseg2"
@@ -83,14 +62,14 @@ def mask_image_api():
     image = Image.open(BytesIO(image_data))
 
     # Process the image
-    _, _, output_image = process_image(image, prompt, threshold, alpha_value, draw_rectangles)
+    output_mask = process_image(image, prompt, threshold, alpha_value, draw_rectangles)
 
-    # Convert the output image to base64
-    buffered = BytesIO()
-    output_image.save(buffered, format="PNG")
-    result_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    # Convert the output mask to base64
+    buffered_mask = BytesIO()
+    output_mask.save(buffered_mask, format="PNG")
+    result_mask = base64.b64encode(buffered_mask.getvalue()).decode('utf-8')
 
-    return jsonify({'result_image': result_image})
+    return jsonify({'result_mask': result_mask})
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860, debug=True)
+    app.run(host='0.0.0.0', port=7860, debug=False)