Initial commit for TorchServe Docker setup

Dockerfile

@@ -0,0 +1,28 @@
+# 1. Pinned Base Image
+FROM pytorch/torchserve:0.12.0-cpu
+
+USER root 
+
+# 2. Install gdown and immediately uninstall it for a smaller image
+RUN pip install --no-cache-dir gdown && pip uninstall gdown -y
+
+# 3. Install other dependencies and clean up
+COPY requirements.txt /tmp/requirements.txt
+RUN pip install -r /tmp/requirements.txt && rm /tmp/requirements.txt
+
+WORKDIR /home/model-server
+RUN mkdir model_store
+
+# 4. Download models using gdown
+RUN gdown 1x9lGntRiYsNb-dYf1SugGGwnoGa_oQes -O model_store/UNET.mar
+RUN gdown 1Y_P77RtNnC1StUeBGlONuNGAW5rFq02S -O model_store/R-UNET.mar
+RUN gdown 1VYvAh5S5MQICbqmQkNJ1Epdmcm5VgVWb -O model_store/A-R-UNET.mar
+
+USER model-server
+
+# 5. Start TorchServe with all models
+CMD ["torchserve", \
+     "--start", \
+     "--ncs", \
+     "--model-store", "/home/model-server/model_store", \
+     "--models", "model_unet=UNET.mar,model_runet=R-UNET.mar,model_arunet=A-R-UNET.mar"]

image_handler.py

@@ -0,0 +1,31 @@
+from ts.torch_handler.base_handler import BaseHandler
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import io
+
+class ImageHandler(BaseHandler):
+    def __init__(self):
+        super(ImageHandler, self).__init__()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.transform = transforms.Compose([transforms.ToTensor()])
+
+    def preprocess(self, data):
+        # TorchServe sends input as bytes → we decode into PIL image
+        image_bytes = data[0].get("body")
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        tensor = self.transform(image).unsqueeze(0).to(self.device)
+        return tensor
+
+    def inference(self, data, *args, **kwargs):
+        with torch.no_grad():
+            output = self.model(data)
+        return output
+
+    def postprocess(self, data):
+        output_tensor = data.squeeze(0).cpu().clamp(0, 1)  # ensure valid range
+        output_image = transforms.ToPILImage()(output_tensor)
+
+        buf = io.BytesIO()
+        output_image.save(buf, format="PNG")
+        return [buf.getvalue()]

requirements.txt

@@ -0,0 +1,4 @@
+torch==2.8.0+cu126
+torchvision==0.23.0+cu126
+numpy==2.0.2
+pillow==11.3.0