Update app.py

app.py

@@ -31,7 +31,7 @@ async def predict(file: UploadFile = File(...)):
     logits = outputs.logits  # (batch, num_labels, H, W)
     mask = torch.argmax(logits, dim=1)[0].numpy().astype(np.uint8)
 
-    # Optionally, you can convert mask to RGB with a color map for visualization
+    # Convert mask to grayscale PNG and return as base64
     mask_img = Image.fromarray(mask)
     buf = BytesIO()
     mask_img.save(buf, format="PNG")