update handler

.gitattributes

@@ -31,3 +31,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+elevator.jpg filter=lfs diff=lfs merge=lfs -text

handler.py

@@ -23,7 +23,6 @@ class EndpointHandler():
       Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-
         if "image" not in data or "text" not in data:
             return [{"error": "Missing 'image' or 'text' key in input data"}]
 

test.ipynb

@@ -0,0 +1,104 @@
+{
+ "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
+}