Update color mask app.py

app.py

@@ -61,8 +61,10 @@ async def predict(image: UploadFile = File(...)):
         logits = F.interpolate(logits, size=(256, 256), mode="bilinear", align_corners=False)
         pred_mask = logits.argmax(dim=1).squeeze().cpu().numpy()
 
-    mask_rgb = decode_mask(pred_mask)
-    mask_img = Image.fromarray(mask_rgb)
+    #mask_rgb = decode_mask(pred_mask)
+    #mask_img = Image.fromarray(mask_rgb)
+    mask_img = Image.fromarray(pred_mask.astype(np.uint8))  # Masque brut
+
 
     buf = io.BytesIO()
     mask_img.save(buf, format="PNG")