Upload app.py

app.py

@@ -40,13 +40,17 @@ async def segment(req: SegmentRequest):
         inputs["original_sizes"].cpu(),
         inputs["reshaped_input_sizes"].cpu()
     )
-    mask_tensor = masks[0][0]                # torch.FloatTensor
-    mask_np     = (mask_tensor * 255)        \
-                    .to(torch.uint8)       \
-                    .cpu()                 \
-                    .numpy()               # now uint8 numpy array
+    single_mask = masks[0][0]                    # first image, first mask → torch.Tensor of shape (H, W)
+    mask_np     = (single_mask * 255)            \
+                    .to(torch.uint8)           \
+                    .cpu().numpy()             
+
+    # now mask_np.shape == (H, W), e.g. (10, 10)
+    print(mask_np.shape)
+    pil_mask = Image.fromarray(mask_np[0,:,:])
     buf = io.BytesIO()
-    Image.fromarray(mask_np).save(buf, format="PNG")
+    pil_mask.save(buf, format="PNG")
     buf.seek(0)
 
+
     return StreamingResponse(buf, media_type="image/png")