load model

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.10-slim
+
+WORKDIR /code
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,58 @@
+from fastapi import FastAPI, UploadFile, File
+import torch
+from safetensors.torch import load_file
+from torchvision import transforms
+from PIL import Image
+import io
+from torchvision import models
+app = FastAPI()
+
+# =============== LOAD MODEL SAFETENSORS ==================
+class Dense121(nn.Module):
+    def __init__(self, num_classes, pretrained=True):
+        super(Dense121, self).__init__()
+        if pretrained:
+            try:
+                weights = models.DenseNet121_Weights.IMAGENET1K_V1
+                self.dense121 = models.densenet121(weights=weights)
+            except:
+                self.dense121 = models.densenet121(pretrained=True)
+        else:
+            self.dense121 = models.densenet121(pretrained=False)
+
+        # classifier DenseNet bukan list, jadi langsung akses in_features
+        in_features = self.dense121.classifier.in_features
+        self.dense121.classifier = nn.Linear(in_features, num_classes)
+
+    def forward(self, x):
+        return self.dense121(x)
+
+
+model = Dense121()
+state_dict = load_file("model_stunting.safetensors")
+model.load_state_dict(state_dict)
+model.eval()
+
+# =============== IMAGE PREPROCESS ==================
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),             # [0..1]
+])
+
+# =============== API ENDPOINT ==================
+@app.post("/predict-image")
+async def predict(file: UploadFile = File(...)):
+
+    img_bytes = await file.read()
+    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    tensor = preprocess(img).unsqueeze(0)   # shape: (1, 3, 224, 224)
+
+    with torch.no_grad():
+        output = model(tensor)
+        probs = torch.softmax(output, dim=1).numpy().tolist()[0]
+
+    return {
+        "prediction": probs,
+        "label": "stunting" if probs[1] > 0.5 else "normal",
+        "confidence": max(probs),
+    }

requirements.txt

@@ -0,0 +1,8 @@
+fastapi
+uvicorn
+python-multipart
+torch
+torchvision
+safetensors
+pillow
+numpy