update handler

handler.py

@@ -56,7 +56,9 @@ class EndpointHandler():
             
             seg_image = Image.fromarray(segmentation_mask)
 
-            return [{"seg_image": seg_image}]
+            seg_image_base64 = self.encode_image(seg_image)
+
+            return [{"seg_image": seg_image_base64}]
         
         except Exception as e:
             return [{"error": str(e)}]
@@ -66,6 +68,12 @@ class EndpointHandler():
         """Decodes a base64-encoded image into a PIL image."""
         image_bytes = base64.b64decode(image_data)
         return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+    def encode_image(self, image: Image.Image) -> str:
+        """Encodes a PIL image to a base64 string."""
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
     
     def process_depth(self, image):
         print("Processing depth")

test.ipynb

@@ -1,104 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Device set to use cpu\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[\n",
-      "  {\n",
-      "    \"error\": \"Torch not compiled with CUDA enabled\"\n",
-      "  }\n",
-      "]\n"
-     ]
-    }
-   ],
-   "source": [
-    "import base64\n",
-    "import json\n",
-    "import io\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
-    "from handler import EndpointHandler \n",
-    "\n",
-    "def encode_image(image_path):\n",
-    "    \"\"\"Encodes an image to base64.\"\"\"\n",
-    "    with open(image_path, \"rb\") as image_file:\n",
-    "        return base64.b64encode(image_file.read()).decode(\"utf-8\")\n",
-    "\n",
-    "def visualize_mask(mask):\n",
-    "    \"\"\"Visualizes the segmentation mask using matplotlib.\"\"\"\n",
-    "    mask = np.array(mask)\n",
-    "    plt.imshow(mask, cmap=\"jet\", alpha=0.6)\n",
-    "    plt.colorbar()\n",
-    "    plt.title(\"Segmentation Mask\")\n",
-    "    plt.show()\n",
-    "\n",
-    "if __name__ == \"__main__\":\n",
-    "    # Initialize handler\n",
-    "    handler = EndpointHandler()\n",
-    "\n",
-    "    # Provide path to a test image\n",
-    "    image_path = \"elevator.jpg\"  # Replace with a valid image path\n",
-    "\n",
-    "    # Encode image as base64\n",
-    "    encoded_image = encode_image(image_path)\n",
-    "\n",
-    "    # Define test input\n",
-    "    test_input = {\n",
-    "        \"image\": encoded_image,\n",
-    "        \"text\": \"elevator\"\n",
-    "    }\n",
-    "\n",
-    "    # Run handler\n",
-    "    result = handler(test_input)\n",
-    "\n",
-    "    # Print and visualize the result\n",
-    "    print(json.dumps(result, indent=2))\n",
-    "\n",
-    "    if \"segmentation_mask\" in result[0]:\n",
-    "        visualize_mask(result[0][\"segmentation_mask\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "a11y-agent",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}