Update app.py

app.py

@@ -18,7 +18,10 @@ def inference(repo_id, model_name, img):
     pred = argmax.cpu().numpy()
     std = stdeviation.cpu().numpy()
     
-    return pred*255, std
+    # Adjust mask for binary cases
+    if model.num_classes==2: pred *= 255
+    
+    return pred, std
  
     
 title="deepflash2"