Update app.py

app.py

@@ -14,22 +14,9 @@ st.set_page_config(layout="wide")
 st.title("Autolamella Demo")
 st.write("This is a demo space of the Autolamella models. Upload FIBSEM (tiff) image, and select a model to segment them.")
 
-st.header(f"Models")
-st.write("""The following models are available for inference. They are trained on different datasets, and may perform differently on different samples.
-         \nautolamella-waffle* is trained on waffle method data.
-         \nautolamella-serial-liftout* is trained on serial liftout data.
-         \nautolamella-mega* is trained on a combination of waffle, autoliftout and serial liftout data. 
-         If you have a new sample, try all of them and see which one works best.""")
-
 # get filenames
 filenames = st.sidebar.file_uploader("Upload an image", type=["tiff", "tif"], accept_multiple_files=True)
 
-# get default data from path
-if len(filenames) == 0:
-    st.write("No files uploaded, using default data")
-    # get all tiff files in current directory
-    filenames = sorted(glob.glob("*.tif"))
-
 # get model
 checkpoint = st.sidebar.selectbox(
     "Select a model checkpoint",
@@ -37,8 +24,24 @@ checkpoint = st.sidebar.selectbox(
      "autolamella-mega-20240107.pt",
      "autolamella-waffle-20240107.pt", 
      "autolamella-serial-liftout-20240107.pt",
-     ],
-)
+     ],)
+
+st.sidebar.header("Available Models")
+st.sidebar.write("""The following models are available for inference. They are trained on different datasets, and may perform differently on different samples.""")
+
+# write a markdown list, listing each of the models
+st.sidebar.write("""
+* autolamella-waffle*
+* autolamella-serial-liftout*
+* autolamella-mega*""")         
+st.sidebar.write("If you have a new sample, try all of them and see which one works best.""")
+
+# get default data from path
+if len(filenames) == 0:
+    st.write("No files uploaded, using default data")
+    # get all tiff files in current directory
+    filenames = sorted(glob.glob("example/*.tif"))
+
 
 st.header(f"Segmentation Results")
 
@@ -46,34 +49,27 @@ st.header(f"Segmentation Results")
 model = load_model(checkpoint)
 
 if filenames:
-    # if there are multiple files, show them in a grid of 4 images
-    if len(filenames) > 1:
-        # create a grid of 4 images
-        n_rows = len(filenames) // 4 + 1
-        if len(filenames) % 4 == 0:
-            n_rows -= 1
-        fig, ax = plt.subplots(n_rows, 4, figsize=(20, 4 * n_rows))
-        ax = ax.flatten()
-        for i, fname in enumerate(filenames):
-
-            # load image, segment, and save
-            image = FibsemImage.load(fname)
-
-            # resize to 1536 x 1024
-            image.data = np.asarray(Image.fromarray(image.data).resize((1536, 1024)))
     
-            mask = model.inference(image.data, rgb=True)
+    cols = st.columns(4)
 
-            # plot
-            ax[i].imshow(image.data, cmap="gray")
-            ax[i].imshow(mask, alpha=0.5)
-            # ax[i].set_title(fname.name)
-            ax[i].axis("off")
-        
-        # remove extra axes
-        for i in range(len(filenames), len(ax)):
-            ax[i].axis("off")
+    for i, fname in enumerate(filenames):
         
-        plt.subplots_adjust(wspace=0, hspace=0)
-        st.pyplot(fig, use_container_width=True)
+        col_id = i % 4
+        cols[col_id].write(f"#### File: {fname} ({i+1}/{len(filenames)})")
+        cols[col_id].write(f"Running inference... ")
+        # load image, segment, and save
+        image = FibsemImage.load(fname)
+
+        # resize to 1536 x 1024
+        image.data = np.asarray(Image.fromarray(image.data).resize((1536, 1024)))
+
+        mask = model.inference(image.data, rgb=True)
+
+        # plot
+        fig = plt.figure(figsize=(10, 10))
+        plt.imshow(image.data, cmap="gray")
+        plt.imshow(mask, alpha=0.5)
+        plt.axis("off")
+
 
+        cols[col_id].pyplot(fig, use_container_width=True)