Update app.py

app.py

@@ -99,12 +99,12 @@ class model:
             
             
         nucleus_image = image['image'].convert('L')
-        protein_image = image['mask'].split()[3]
+        #protein_image = image['mask'].split()[3]
+        protein_image = image['mask'].convert('L')
 
         to_tensor = T.ToTensor()
         nucleus_image = to_tensor(nucleus_image)
         protein_image = to_tensor(protein_image)
-        #protein_image = protein_image
         #stacked_images = torch.stack([nucleus_image, protein_image], dim=0)
         #processed_images = process_image(stacked_images, dataset)
 
@@ -112,10 +112,12 @@ class model:
         #protein_image = processed_images[1].unsqueeze(0)
         #protein_image = protein_image/torch.max(protein_image)
         #protein_image = 1 - protein_image  
-
+        
         nucleus_image = nucleus_image.unsqueeze(0)
         nucleus_image = process_image(nucleus_image, dataset, 'nucleus')
-        protein_image = 1.0*(protein_image > .01)
+        nucleus_image = nucleus_image.unsqueeze(0)
+        protein_image = protein_image.unsqueeze(0).unsqueeze(0)
+        #protein_image = 1.0*(protein_image > .01)
         
 
         formatted_predicted_sequence = run_sequence_prediction(
@@ -168,13 +170,19 @@ with gr.Blocks(theme='gradio/soft') as demo:
         )
 
     with gr.Row().style(equal_height=True):
-        nucleus_image = gr.Image(
-            source="upload", 
-            tool="color-sketch",
-            label="Nucleus Image", 
-            interactive=True,
-            image_mode="RGBA",
-            type="pil"
+        #nucleus_image = gr.Image(
+        #    source="upload", 
+        #    tool="color-sketch",
+        #    label="Nucleus Image", 
+        #    interactive=True,
+        #    image_mode="RGBA",
+        #    type="pil"
+        #)
+        nucleus_image = gr.ImageMask(
+            label = "Nucleus Image",
+            interactive = "True",
+            image_mode = "L",
+            brush_color = "#FFFFFF"
         )
         
     with gr.Row():