Merge pull request #641 from czbiohub-sf/multi-head-classification-ssaf

Multi head classification ssaf

Author: i-jey

ulc_mm_package/QtGUI/oracle.py

@@ -5,13 +5,14 @@
 
 """
 
-import os
-import sys
-import traceback
-import socket
 import enum
 import logging
+import os
 import subprocess
+import socket
+import sys
+import traceback
+from typing import Optional
 
 from os import (
     listdir,
@@ -70,11 +71,13 @@
 from ulc_mm_package.neural_nets.neural_network_constants import (
     AUTOFOCUS_MODEL_DIR,
     YOGO_MODEL_DIR,
+    QC_STATUS,
 )
 
 from ulc_mm_package.QtGUI.scope_op import ScopeOp
 from ulc_mm_package.QtGUI.form_gui import FormGUI
 from ulc_mm_package.QtGUI.liveview_gui import LiveviewGUI
+from PyQt5.QtWidgets import QPushButton
 
 
 QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
@@ -762,11 +765,86 @@ def _start_liveview(self, *args):
     def _end_liveview(self, *args):
         self.liveview_window.close()
 
-    def _start_intermission(self, msg=None, parasitemia_vis_path=""):
+    def _start_intermission(
+        self,
+        msg=None,
+        parasitemia_vis_path="",
+        run_qc_status: Optional[int] = None,
+    ):
         if msg is None:
             # Retriggered intermission due to race condition
             return
 
+        # Display QC results if available
+        # An Enum would be great but `pyqtsignal` on PyQt5 does not support Enums
+        # and a workaround would be uglier
+        if run_qc_status is not None:
+            subsample_dir = self.scopeop.mscope.data_storage.get_subsample_folder_path()
+            if not os.path.exists(subsample_dir):
+                self.logger.error(
+                    f"Subsample images directory does not exist: {subsample_dir}"
+                )
+                subsample_dir = None
+            if run_qc_status == QC_STATUS.GOOD.value:
+                # Add a custom button labeled
+                investigate_btn = None
+
+                msg_box = NoCloseMessageBox()
+                msg_box.setWindowIcon(QIcon(ICON_PATH))
+                msg_box.setIcon(QMessageBox.Icon.Information)
+                msg_box.setWindowTitle("Run Quality: GOOD")
+                msg_box.setText("✅ The run quality is GOOD.\n\n")
+
+                # Ok button
+                msg_box.addButton(QMessageBox.Ok)
+
+                # Add button to open subsample images
+                if subsample_dir is not None:
+                    investigate_btn = QPushButton("View subsample images")
+                    msg_box.addButton(investigate_btn, QMessageBox.ActionRole)
+                    msg_box.exec()
+                    if msg_box.clickedButton() == investigate_btn:
+                        # Path to the subsample images directory
+                        subsample_dir = (
+                            self.scopeop.mscope.data_storage.get_subsample_folder_path()
+                        )
+                        if os.name == "posix":
+                            subprocess.call(["xdg-open", subsample_dir])
+
+            elif run_qc_status == QC_STATUS.POOR.value:
+                investigate_btn = None
+
+                msg_box = NoCloseMessageBox()
+                msg_box.setWindowIcon(QIcon(ICON_PATH))
+                msg_box.setIcon(QMessageBox.Icon.Critical)
+                msg_box.setWindowTitle("Run Quality: POOR")
+                msg_box.setText(
+                    "❌ The run quality is POOR.\n\nPlease RE-RUN THE SAMPLE to ensure valid results."
+                )
+                msg_box.setDetailedText(
+                    (
+                        "Open the run's subsample images to investigate its quality. The QC result is usually indicative of image quality, however "
+                        "in cases where the images are unusual for another reason (anemia, SCD, other hemoglobinopathies), the QC result may not be accurate."
+                        "If the images appear in-focus (i.e NOT blurry), and the cells are NOT coagulated, then you do not need to re-run the sample. If you are unsure, please re-run the sample."
+                    )
+                )
+
+                # Ok button
+                msg_box.addButton(QMessageBox.Ok)
+
+                # Add button to open subsample images
+                if subsample_dir is not None:
+                    investigate_btn = QPushButton("View subsample images")
+                    msg_box.addButton(investigate_btn, QMessageBox.ActionRole)
+                    msg_box.exec()
+                    if msg_box.clickedButton() == investigate_btn:
+                        if os.name == "posix":
+                            subprocess.call(["xdg-open", subsample_dir])
+            else:
+                raise ValueError(
+                    f"Invalid run_qc_status: {run_qc_status}. Expected 'good' or 'poor'."
+                )
+
         self.display_message(
             QMessageBox.Icon.Information,
             "Run status",

ulc_mm_package/QtGUI/scope_op.py

@@ -5,15 +5,15 @@
 
 """
 
-import cv2
 import logging
-import numpy as np
 
 from typing import Any
 from time import sleep, perf_counter
-from transitions import Machine, State
 
+import cv2
+import numpy as np
 from PyQt5.QtCore import QObject, pyqtSignal, pyqtSlot
+from transitions import Machine, State
 
 from ulc_mm_package.hardware.scope import MalariaScope, GPIOEdge
 from ulc_mm_package.hardware.scope_routines import Routines
@@ -38,7 +38,12 @@
     LEDNoPower,
 )
 
-from ulc_mm_package.neural_nets.neural_network_constants import IMG_RESIZED_DIMS
+from ulc_mm_package.neural_nets.neural_network_constants import (
+    IMG_RESIZED_DIMS,
+    QC_GOODNESS_THRESHOLD,
+    QC_STATUS,
+    PERC_OF_IMAGES_GOOD,
+)
 from ulc_mm_package.neural_nets.YOGOInference import YOGO, ClassCountResult
 from ulc_mm_package.neural_nets.neural_network_constants import (
     YOGO_CLASS_LIST,
@@ -84,7 +89,7 @@ class NamedMachine(Machine):
 
 class ScopeOp(QObject, NamedMachine):
     setup_done = pyqtSignal()
-    experiment_done = pyqtSignal(str, str)
+    experiment_done = pyqtSignal(str, str, int)
     reset_done = pyqtSignal()
 
     yield_mscope = pyqtSignal(MalariaScope)
@@ -301,6 +306,34 @@ def update_thumbnails(self):
                 )
             )
 
+    def run_status_from_qc_results(self, qc_results: np.ndarray) -> QC_STATUS:
+        """Logic for determining whether a run was decent based on the QC results at the end of a run.
+
+        Parameters
+        ----------
+        qc_results : np.ndarray
+            Array of QC results from the QC model.
+
+        Returns
+        -------
+        QC_STATUS
+            Status of the run based on the QC results.
+            - "good" if X% of results are below the QC_GOODNESS_THRESHOLD (i.e considered 'good')
+            - "poor" if Y% of results are above the threshold
+        """
+
+        num_good = (qc_results <= QC_GOODNESS_THRESHOLD).sum()
+        num_total = len(qc_results)
+
+        if num_total == 0:
+            self.logger.warning("No QC results available. Cannot determine run status.")
+            raise ValueError("Run status cannot be determined without QC results.")
+
+        if num_good / num_total >= PERC_OF_IMAGES_GOOD:
+            return QC_STATUS.GOOD
+        else:
+            return QC_STATUS.POOR
+
     def setup(self):
         self.create_timers.emit()
 
@@ -559,6 +592,48 @@ def _end_experiment(self, *args):
 
         self.mscope.reset_for_end_experiment()
 
+        # Run the QC model a small partition of the data
+        self.logger.info("Running QC on images.")
+        qc_results = None
+        try:
+            zf = self.mscope.data_storage.get_read_only_zarr()
+
+            try:
+                img_indices = np.linspace(0, zf.initialized - 1, 50).astype(int)
+                for idx in img_indices:
+                    img = zf[:, :, idx]
+                    self.mscope.qc.asyn(img)
+                qc_results = self.mscope.qc.get_asyn_results(timeout=None)
+                qc_results = [self.mscope.qc._sigmoid(x.result) for x in qc_results]
+            except Exception as e:
+                self.logger.error(
+                    f"Unexpected error while submitting images to QC model: {e}. Skipping QC..."
+                )
+        except Exception as e:
+            self.logger.error(f"Failed to get zarr data for QC: {e}.\nSkipping QC...")
+
+        # Log QC results
+        if qc_results:
+            self.did_run_pass_qc = None
+            qc_results_np = np.array([x[0][0] for x in qc_results])
+            num_qc_results_good = (qc_results_np <= QC_GOODNESS_THRESHOLD).sum()
+            num_imgs_qc = len(qc_results_np)
+            self.logger.info(
+                f"QC all results: {qc_results_np}\n"
+                f"QC mean: {qc_results_np.mean():.3f}, "
+                f"QC stdev: {qc_results_np.std():.3f}, "
+                f"QC best score: {qc_results_np.min():.3f}, "
+                f"QC worst score: {qc_results_np.max():.3f}, "
+                f"QC number of images good: {num_qc_results_good}/{num_imgs_qc} ({num_qc_results_good/num_imgs_qc:.2%})%"
+            )
+            self.did_run_pass_qc = self.run_status_from_qc_results(qc_results_np).value
+        else:
+            self.logger.warning("No QC results available. Skipping QC...")
+
+        # Save qc results
+        self.mscope.data_storage.save_qc_data(img_indices, qc_results_np)
+        self.finishing_experiment.emit(80)
+
         # Turn camera back on
         self.mscope.camera.startAcquisition()
 
@@ -577,12 +652,15 @@ def _end_experiment(self, *args):
     def _start_intermission(self, msg):
         parasitemia_vis_path = self.mscope.data_storage.get_parasitemia_vis_filename()
 
+        # Display parasitemia visualization if it exists
         if parasitemia_vis_path.exists():
             self.experiment_done.emit(
-                msg + PARASITEMIA_VIS_MSG, str(parasitemia_vis_path)
+                msg + PARASITEMIA_VIS_MSG,
+                str(parasitemia_vis_path),
+                self.did_run_pass_qc,
             )
         else:
-            self.experiment_done.emit(msg, "")
+            self.experiment_done.emit(msg, "", self.did_run_pass_qc)
 
     @pyqtSlot(np.ndarray, float)
     def run_autobrightness(self, img, _timestamp):
@@ -701,10 +779,7 @@ def run_autofocus(self, img, _timestamp):
 
         if not self.autofocus_done:
             if len(self.autofocus_batch) < AF_BATCH_SIZE:
-                resized_img = cv2.resize(
-                    img, IMG_RESIZED_DIMS, interpolation=cv2.INTER_CUBIC
-                )
-                self.autofocus_batch.append(resized_img)
+                self.autofocus_batch.append(img)
 
                 if self.running:
                     self.img_signal.connect(self.run_autofocus)
@@ -931,7 +1006,7 @@ def run_experiment(self, img, timestamp) -> None:
                 raw_focus_err,
                 filtered_focus_err,
                 focus_adjustment,
-            ) = self.PSSAF_routine.send(resized_img)
+            ) = self.PSSAF_routine.send(img)
         except MotorControllerError as e:
             if not SIMULATION:
                 self.logger.error(

ulc_mm_package/hardware/scope.py

@@ -1,6 +1,6 @@
 """
 The malaria scope object, containing all the different hardware
-periperhals which make up the malaria scope.
+peripherals which make up the malaria scope.
 
 Components
     - Motorcontroller (DRV8825 Nema)
@@ -50,6 +50,7 @@
 from ulc_mm_package.image_processing.flow_control import FlowController
 from ulc_mm_package.neural_nets.YOGOInference import YOGO
 from ulc_mm_package.neural_nets.AutofocusInference import AutoFocus
+from ulc_mm_package.neural_nets.QCInference import QC
 from ulc_mm_package.neural_nets.NCSModel import GPUError
 from ulc_mm_package.neural_nets.predictions_handler import PredictionsHandler
 
@@ -129,7 +130,7 @@ def reset_pneumatic_and_led_and_flow_control(self) -> None:
 
     def reset_for_end_experiment(self) -> None:
         """
-        Reset syringe, turn LED off, reset flow control, reset YOGO / Autofoucs,
+        Reset syringe, turn LED off, reset flow control, reset YOGO / Autofocus,
         and close data storage.
         """
 
@@ -322,6 +323,7 @@ def _init_GPU(self):
             self.logger.info("Initializing GPU...")
             self.autofocus_model = AutoFocus()
             self.cell_diagnosis_model = YOGO()
+            self.qc = QC()
             self.gpu_enabled = True
         except GPUError as e:
             self.logger.error(f"GPU initialization failed. {e}")

ulc_mm_package/hardware/scope_routines.py

@@ -76,8 +76,16 @@ def singleShotAutofocusRoutine(
             The number of steps that the motor was moved.
         """
 
-        ssaf_steps_from_focus = mscope.autofocus_model(img_arr)
-        steps_to_move = -round(np.mean(ssaf_steps_from_focus))
+        res = mscope.autofocus_model(img_arr)
+
+        # The indices of the first output correspond to below/in/above focus.
+        # We subtract one so that it maps to -1, 0, +1, and then subsequently take the element-wise product with the
+        # magnitude vector to get the number of steps away from focus.
+        above_or_below_vec = np.array([np.argmax(r[0][0]) - 1 for r in res])
+        magnitude_vec = np.array([np.argmax(r[1][0]) for r in res])
+        steps_vec = above_or_below_vec * magnitude_vec
+
+        steps_to_move = -round(np.mean(steps_vec))
 
         try:
             dir = Direction.CW if steps_to_move > 0 else Direction.CCW
@@ -147,11 +155,21 @@ def periodicAutofocusWrapper(
                 mscope.autofocus_model.asyn(img, img_counter)
                 results = mscope.autofocus_model.get_asyn_results(timeout=0.005) or []
 
-                for res in sorted(results, key=lambda res: res.id):
+                for res in sorted(results, key=lambda res: res[0].id):
                     move_counter += 1
-
-                    steps_from_focus = res.result.item()
-                    filtered_error = ssaf_filter.update_and_get_val(steps_from_focus)
+                    direction_asyn_result = res[0]
+                    mag_asyn_result = res[1]
+
+                    direction = (
+                        np.argmax(direction_asyn_result.result) - 1
+                    )  # Subtract 1 so that it maps to -1, 0, +1
+                    magnitude = np.argmax(mag_asyn_result.result)
+                    mag_conf = mag_asyn_result.result[0][magnitude]
+                    if mag_conf >= nn_constants.MAG_CONF_THRESH:
+                        steps_from_focus = direction * magnitude
+                        filtered_error = ssaf_filter.update_and_get_val(
+                            steps_from_focus
+                        )
 
                 throttle_counter = 0