change samples

app.py

@@ -47,7 +47,7 @@ pe_model = Cholec80Model({"image": [2048, 128], "pos_enc": [7, 7, 128]})
 load_pretrained_params(pe_model, "checkpoints/cnn_pe_2.ckpt")
 
 
-def cnn(image):
+def cnn(label, image):
     # unsqueeze the input_tensor
     input_tensor = transform(image)
     input_tensor = input_tensor.unsqueeze(dim=0).to(device)
@@ -60,7 +60,7 @@ def cnn(image):
     return {k: float(pred_softmax[v]) for k, v in classes.items()}
 
 
-def cnn_mask(image, last_phase):
+def cnn_mask(label, image, last_phase):
     # extract last phase
     last_phase = int(last_phase.split("-")[0].strip())
     # mask
@@ -85,7 +85,7 @@ def cnn_mask(image, last_phase):
     return {k: float(pred_softmax[v]) for k, v in classes.items()}
 
 
-def cnn_pe(image, p_0, p_1, p_2, p_3, p_4, p_5, p_6):
+def cnn_pe(label, image, p_0, p_1, p_2, p_3, p_4, p_5, p_6):
     # form the position encoder vector
     pos_enc = torch.Tensor([[p_0, p_1, p_2, p_3, p_4, p_5, p_6]]).to(device)
     # unsqueeze the input_tensor
@@ -103,16 +103,18 @@ with gr.Blocks() as demo:
     gr.Markdown("# Phase Recognition of Cholecystectomy Surgeries")
     # inputs
     with gr.Row():
-        image_input = gr.Image(shape=(255, 255), type="pil")
+        with gr.Column():
+            groundtruth_lable = gr.Text(label="Ground Truth", interactive=False)
+            image_input = gr.Image(shape=(255, 255), type="pil")
         # output
-        lable_output = gr.Label()
+        lable_output = gr.Label(label="Result")
     with gr.Tab("CNN") as cnn_tab:
         cnn_button = gr.Button("Predict")
-        cnn_button.click(cnn, inputs=[image_input], outputs=[lable_output])
+        cnn_button.click(cnn, inputs=[groundtruth_lable, image_input], outputs=[lable_output])
     with gr.Tab("CNN+Mask") as mask_tab:
         phase = gr.Dropdown([f"{v} - {k}" for k, v in classes.items()], label="Last frame is of phase")
         mask_button = gr.Button("Predict")
-        mask_button.click(cnn_mask, inputs=[image_input, phase], outputs=[lable_output])
+        mask_button.click(cnn_mask, inputs=[groundtruth_lable, image_input, phase], outputs=[lable_output])
     with gr.Tab("CNN+PE") as pe_tab:
         with gr.Row():
             p0 = gr.Number(label="Phase 0")
@@ -123,16 +125,16 @@ with gr.Blocks() as demo:
             p5 = gr.Number(label="Phase 5")
             p6 = gr.Number(label="Phase 6")
         pe_button = gr.Button("Predict")
-        pe_button.click(cnn_pe, inputs=[image_input, p0, p1, p2, p3, p4, p5, p6], outputs=[lable_output])
+        pe_button.click(cnn_pe, inputs=[groundtruth_lable, image_input, p0, p1, p2, p3, p4, p5, p6], outputs=[lable_output])
     gr.Examples(
-        examples=[['images/preparation.png'],
-                  ['images/calot-triangle-dissection.png'],
-                  ['images/clipping-cutting.png'],
-                  ['images/gallbladder-dissection.png'],
-                  ['images/gallbladder-packaging.png'],
-                  ['images/cleaning-coagulation.png'],
-                  ['images/gallbladder-retraction.png']],
-        inputs=image_input
+        examples=[['0 - Preparation', 'images/pe-sample.png'],
+                  ['1 - Calot Triangle Dissection', 'images/mask-sample.png'],
+                  ['2 - Clipping Cutting', 'images/clipping-cutting.png'],
+                  ['3 - Gallbladder Dissection', 'images/gallbladder-dissection.png'],
+                  ['4 - Gallbladder Packaging', 'images/gallbladder-packaging.png'],
+                  ['5 - Cleaning Coagulation', 'images/cleaning-coagulation.png'],
+                  ['6 - Gallbladder Retraction', 'images/gallbladder-retraction.png']],
+        inputs=[groundtruth_lable, image_input]
     )
 
 if __name__ == "__main__":