Add application file

app.py

@@ -0,0 +1,69 @@
+import base64
+from io import BytesIO
+import gradio as gr
+from PIL import Image
+import json
+
+from tools.tools import convertToBuffer
+from visualize.visualize import removeBgFromSegmentImage, removeOnlyBg
+from models.model import getMask, loadModel
+from models.preprocess import preprocess
+
+FAST_SAM = loadModel()
+
+def base64_to_image(base64_str):
+    image_data = base64.b64decode(base64_str)
+    image = Image.open(BytesIO(image_data))
+    return image
+
+# Main processing function
+def segment_marker(img_rgb: Image.Image, marker_coordinates: str):
+    # Parse marker coordinates from JSON string
+    try:
+        marker_coordinates = json.loads(marker_coordinates)
+    except json.JSONDecodeError:
+        return "Invalid marker coordinates format. Ensure it's valid JSON."
+
+    try:
+        # Process marker points and labels
+        input_points, input_labels = preprocess(marker_coordinates)
+
+        print(f"Processing image with {len(input_points)} marker points...")
+        # Get mask for segmentation
+        masks = getMask(img_rgb, FAST_SAM, input_points, input_labels)
+
+        # Generate the segmented images
+        bg_removed_segmented_img = removeBgFromSegmentImage(img_rgb, masks[0])
+        img_base64_bg_segmented = convertToBuffer(bg_removed_segmented_img)
+
+        bg_only_removed_img = removeOnlyBg(img_rgb, masks[0])
+        img_base64_only_bg = convertToBuffer(bg_only_removed_img)
+
+        # Convert base64 strings to PIL images for Gradio
+        img_bg_segmented = base64_to_image(img_base64_bg_segmented)
+        img_bg_only_removed = base64_to_image(img_base64_only_bg)
+
+        return img_bg_segmented, img_bg_only_removed  # Return as two separate images
+
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+        return "An error occurred while processing the image.", None
+
+
+# Set up the Gradio interface
+iface = gr.Interface(
+    fn=segment_marker,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Markers Coordinates (JSON format)")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Background Removed with Segmentation"),
+        gr.Image(type="pil", label="Only Background Removed")
+    ],
+    title="Image Segmentation with Background Removal",
+    description="Upload an image and JSON-formatted marker coordinates to perform image segmentation and background removal."
+)
+
+# Run the Gradio app
+iface.launch(share=True)