Update app.py

app.py

@@ -54,17 +54,18 @@ class model:
         
         to_tensor = T.ToTensor()
         
-        nucleus_tensor = to_tensor(nucleus_image)
+        nucleus_image = to_tensor(nucleus_image)
 
         if protein_image:
-            protein_tensor = to_tensor(protein_image)
-            stacked_images = torch.stack([nucleus_tensor, protein_tensor], dim=0)
+            protein_image = to_tensor(protein_image)
+            stacked_images = torch.stack([nucleus_image, protein_image], dim=0)
             processed_images = process_image(stacked_images, dataset)
             nucleus_image = processed_images[0].unsqueeze(0)
             protein_image = processed_images[1].unsqueeze(0)
             protein_image = protein_image > 0
             protein_image = 1.0 * protein_image
         else:
+            nucleus_image = process_image(nucleus_image)
             protein_image = torch.ones((256, 256))
 
         threshold, heatmap = run_image_prediction(