Update app.py

app.py

@@ -51,13 +51,19 @@ class model:
 
         else:
             dataset = "HPA"
+        
+        to_tensor = T.ToTensor()
+        
+        nucleus_tensor = to_tensor(nucleus_image)
 
-        nucleus_image = process_image(nucleus_image, dataset, "nucleus")
         if protein_image:
-            protein_image = process_image(protein_image, dataset, "protein")
-            protein_image = protein_image > torch.median(protein_image)
-            protein_image = protein_image[0, 0]
-            protein_image = protein_image * 1.0
+            protein_tensor = to_tensor(protein_image)
+            stacked_images = torch.stack([nucleus_tensor, protein_tensor], 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:
             protein_image = torch.ones((256, 256))