Add functionality for downloading segmented images

app.py

@@ -1,5 +1,7 @@
 import base64
+import os
 import tempfile
+import zipfile
 from pathlib import Path
 
 import cv2
@@ -35,7 +37,7 @@ main_app = ui.page_fluid(
                               multiple=True,
                               accept=[".png", ".jpg", ".jpeg"]),
 
-                ui.input_slider("threshold", f"Threshold for Discarding Overlapping Segmentations (Default: {OPTIMAL_NMS_THRESHOLD})",
+                ui.input_slider("nms_threshold", f"Threshold for Discarding Overlapping Segmentations (Default: {OPTIMAL_NMS_THRESHOLD})",
                                 0, 1.0,
                                 OPTIMAL_NMS_THRESHOLD),
 
@@ -88,8 +90,9 @@ main_app = ui.page_fluid(
                     """  # Style need adding here for slider for some reason
                               ),
                 ui.input_action_button("analyse", "Analyse", class_="btn-success"),
+                ui.row(class_="analysis-separator"),
                 # ui.input_switch("mask", "Mask", False),
-                ui.output_ui("download_ui"),
+                ui.output_ui("download_results_ui"),
                 width=300
 
             ),
@@ -105,7 +108,7 @@ app_ui = ui.page_fluid(
     ui.div(
         ui.row(
             ui.column(5,
-                      ui.panel_title(ui.div("OrchAid", ui.output_image("logo_image", inline=True, width='100px'), class_="navbar-title"))
+                      ui.panel_title(ui.div("OrchAId", ui.output_image("logo_image", inline=True, width='100px'), class_="navbar-title"))
                       )
         ),
         class_="nav-bar"
@@ -153,8 +156,8 @@ app_ui = ui.page_fluid(
                              " Full details of the model, training process and evaluation can be found on the project ",
                              ui.a("GitHub repository", href=github_repo_url, target="_blank"),
                              ". You can find a project overview ", ui.a("here",
-                                                                           href='https://www.kew.org/science/our-science/projects/machine-learning-to-improve-orchid-viability-testing',
-                                                                           target="_blank"), '.'),
+                                                                        href='https://www.kew.org/science/our-science/projects/machine-learning-to-improve-orchid-viability-testing',
+                                                                        target="_blank"), '.'),
                          class_="body-bar"))
         , id='tab'
     ),
@@ -189,6 +192,28 @@ def plot_ui():
     )
 
 
+def get_overlayed_image_from_single_result(r, opacity=0.5, palette=None):
+    '''
+    From the stored result, get the overlayed image.
+    :param r:
+    :param opacity:
+    :return:
+    '''
+    v = Visualizer(r["image"][:, :, ::-1],
+                   scale=1.2, instance_mode=ColorMode.SEGMENTATION, font_size_scale=1)
+
+    if palette is None:
+        palette = [[1, 0, 0], [1, 1, 0], [0, 0, 0]]
+
+    colours = []
+    for cls in r["instances"].pred_classes:
+        colours.append(palette[cls])
+    out = v.overlay_instances(masks=r["instances"].pred_masks.to("cpu"),
+                              assigned_colors=colours,
+                              alpha=opacity)
+    return out
+
+
 @module.server
 def plot_server(input, output, session, r):
     @render.plot
@@ -200,21 +225,7 @@ def plot_server(input, output, session, r):
         ax.set_axis_off()
         # fig.add_axes(ax)
 
-        v = Visualizer(r["image"][:, :, ::-1],
-                       scale=1.2, instance_mode=ColorMode.SEGMENTATION, font_size_scale=1)
-
-        colours = []
-        for cls in r["instances"].pred_classes:
-            if cls == 0:
-                colours.append([1, 0, 0])
-            elif cls == 1:
-                colours.append([1, 1, 0])
-            elif cls == 2:
-                colours.append([0, 0, 0])
-
-        out = v.overlay_instances(masks=r["instances"].pred_masks.to("cpu"),
-                                  assigned_colors=colours,
-                                  alpha=input.opacity_slider())
+        out = get_overlayed_image_from_single_result(r, input.opacity_slider())
 
         ax.imshow(cv2.cvtColor(out.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB))
         fig.canvas.draw()
@@ -277,14 +288,14 @@ def server(input, output, session: Session):
                 # Run prediction with original BGR image
                 prediction = predictor(im)
                 print(f"Analyzing image {idx + 1} of {len(files)}")
-                print(f"NMS threshold: {input.threshold()}")
+                print(f"NMS threshold: {input.nms_threshold()}")
                 print(f'Number of instances: {len(prediction["instances"])}')
-                prediction = apply_nms(prediction, mask=True, cls_agnostic_nms=input.threshold())
+                prediction = apply_nms(prediction, mask=True, cls_agnostic_nms=input.nms_threshold())
                 print(f'Number of instances after NMS: {len(prediction["instances"])}')
 
                 classes = prediction["instances"].pred_classes.tolist()
 
-                results.append({
+                single_result = {
                     "filename": file["name"],
                     "image_base64": img_base64,
                     "image": im,
@@ -293,8 +304,10 @@ def server(input, output, session: Session):
                     "non-viable": classes.count(1),
                     "empty": classes.count(2),
                     "total": len(classes),
-                    'NMS threshold': input.threshold()
-                })
+                    'NMS threshold': input.nms_threshold()
+                }
+                results.append(single_result)
+                # print(f'Size of result: {sys.getsizeof(single_result)} bytes')
 
         # Update reactive value
         analysis_results.set(results)
@@ -335,12 +348,14 @@ def server(input, output, session: Session):
         return ui.div(ui_output)
 
     @render.ui
-    def download_ui():
+    def download_results_ui():
         if analysis_results.get() and not is_analyzing.get():
-            return ui.column(4, ui.download_button("download", "Download Results", class_="btn-success"))
+            return ui.download_button("download_results", "Download Results", class_="btn-success"), ui.download_button("download_segmented_images",
+                                                                                                                        "Download Segmented Images",
+                                                                                                                        class_="btn-success")
 
     @render.download()
-    def download():
+    def download_results():
         results = analysis_results.get()
         # if not results:
         #     None
@@ -355,10 +370,34 @@ def server(input, output, session: Session):
         } for r in results])
 
         # Create in-memory CSV file
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
+        with tempfile.NamedTemporaryFile(delete=False, delete_on_close=True, suffix=".csv") as tmp:
+            print(f'result tmp csv: {tmp.name}')
             df.to_csv(tmp.name, index=False)
             return tmp.name
 
+    @render.download()
+    def download_segmented_images():
+        results = analysis_results.get()
+
+        tmp_img_files = []
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            print(os.listdir(os.path.dirname(temp_dir)))
+            for r in results:
+                # open your files here
+                named_file = os.path.join(temp_dir, r['filename'])
+                img = get_overlayed_image_from_single_result(r)
+                img.save(named_file)
+                tmp_img_files.append(named_file)
+
+            with tempfile.NamedTemporaryFile(delete=False, delete_on_close=True, suffix=".zip") as tmp:
+
+                with zipfile.ZipFile(tmp.name, 'w') as zipMe:
+                    for file in tmp_img_files:
+                        zipMe.write(file, compress_type=zipfile.ZIP_DEFLATED)
+
+                return tmp.name
+
 
 app = App(app_ui, server)
 

styles.css

@@ -62,6 +62,10 @@ position: sticky;
     border-top: 2px solid #ddd;
     background: white;
 }
+.analysis-separator {
+    border-top: 2px solid #ddd;
+    background: white;
+}
 
 /* Sidebar styling */
 .card.shiny-input-container {