Update app.py

app.py

@@ -26,13 +26,13 @@ def inference(image, chosen_model):
     logits = outputs.logits
 
     output = torch.sigmoid(logits).detach().numpy()[0]
-    output = np.transpose(output, (1,2,0))
-    upsampled_logits = nn.functional.interpolate(logits,
-                size=image.size[::-1], # (height, width)
-                mode='bilinear',
-                align_corners=False)
+    # output = np.transpose(output, (1,2,0))
+    # upsampled_logits = nn.functional.interpolate(logits,
+    #             size=image.size[::-1], # (height, width)
+    #             mode='bilinear',
+    #             align_corners=False)
 
-    seg = upsampled_logits.argmax(dim=1)[0]
+    seg = output
     color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
     palette = np.array([[0, 0, 0],[255, 255, 255]])
     for label, color in enumerate(palette):