From c56d273be4606f3eea34efc1bc70ed208b0191a5 Mon Sep 17 00:00:00 2001 From: Jakub Both Date: Sun, 26 Nov 2023 22:41:14 +0100 Subject: [PATCH] ENH: Add more functionality (monochromatic, background toggle, info, ...) --- src/darsia/assistants/labels_assistant.py | 332 ++++++++++++++++------ 1 file changed, 250 insertions(+), 82 deletions(-) diff --git a/src/darsia/assistants/labels_assistant.py b/src/darsia/assistants/labels_assistant.py index 088e2b47..d8f844e9 100644 --- a/src/darsia/assistants/labels_assistant.py +++ b/src/darsia/assistants/labels_assistant.py @@ -1,25 +1,19 @@ """"Labels assistant built from modules.""" from typing import Optional +from warnings import warn import matplotlib.pyplot as plt import numpy as np import darsia -# TODOs: -# Add possibility to manipulate the monochromatic input image for segmentation -# Test plain segmentation from start -# Add possibility to control the boundary during the segmentation - class LabelsAssistantMenu(darsia.BaseAssistant): """Module for LabelsAssistant to choose what to do with labels. - Purpose is to forward which module to use to LabelsAssistant. Choices: - - Merge labels - - Choose labels - - Segment labels + Purpose is to forward which module to use to LabelsAssistant. Detailed instructions + are printed to screen. """ @@ -37,20 +31,25 @@ def __init__( """Name of the module.""" # Initialize base assistant with reference to current labels - super().__init__(img=img, background=background, block=True, **kwargs) + super().__init__(img=img, background=background, **kwargs) # Print instructions - self._print_instructions() + if kwargs.get("print_instructions", True): + self._print_instructions() def _print_instructions(self) -> None: """Print instructions.""" print("Welcome to the labels assistant.") print("Please choose one of the following options:") - print(" - 's': segment labels") + print(" - 'shift+s': segment labels") print(" - 'm': merge labels") print(" - 'p': pick labels") print(" - 'r': refine labels") + print(" - 'i': print info") + print(" - 'b': toggle background") + print(" - 'shift+m': adapt monochromatic background image") + print(" - 'u': undo") print(" - 'escape': reset labels to input") print(" - 'q': quit/abort\n") @@ -62,32 +61,45 @@ def _on_key_press(self, event) -> None: event: key press event """ - super()._on_key_press(event) + if self.verbosity: + print(f"Current key: {event}") - # Additional events - if event.key == "s": + # Track events + if event.key == "S": self.action = "segment" plt.close(self.fig) - if event.key == "m": + elif event.key == "m": self.action = "merge" plt.close(self.fig) - if event.key == "p": + elif event.key == "p": self.action = "pick" plt.close(self.fig) - if event.key == "r": + elif event.key == "r": self.action = "refine" plt.close(self.fig) - if event.key == "escape": + elif event.key == "escape": self.action = "reset" plt.close(self.fig) - if event.key == "q": + elif event.key == "b": + self.action = "toggle_background" + plt.close(self.fig) + elif event.key == "M": + self.action = "monochromatic" + plt.close(self.fig) + elif event.key == "u": + self.action = "undo" + plt.close(self.fig) + elif event.key == "i": + self.action = "info" + plt.close(self.fig) + elif event.key == "q": self.action = "quit" plt.close(self.fig) def __call__(self) -> darsia.Image: """Call the assistant.""" - self.action = "quit" + self.action = None super().__call__() return self.action @@ -95,69 +107,145 @@ def __call__(self) -> darsia.Image: class LabelsSegmentAssistant: def __init__( self, - labels: darsia.Image, + labels: Optional[darsia.Image], background: darsia.Image, mask: Optional[np.ndarray] = None, **kwargs, ) -> None: - self.labels = labels # If labels provided, ask which to keep + self.labels = labels + """Input labels.""" self.background = background + """Background image.""" + assert self.background is not None, "Background image required." self.mask = mask + """Mask to be used for segmentation.""" self.verbosity = kwargs.get("verbosity", False) + """Verbosity flag.""" + + # Purely for visualization purposes + self.roi = None + """Region of interest.""" # Safety checks if mask is not None: assert isinstance(mask, np.ndarray), "Mask must be a numpy array." assert mask.dtype == bool, "Mask must be a boolean array." + self.roi = darsia.Image( + np.zeros_like( + background.img if background.scalar else background.img[:, :, 0], + dtype=int, + ), + **self.background.metadata(), + ) + self.roi.img[mask] = 1 def __call__(self) -> darsia.Image: point_selection_assistant = darsia.PointSelectionAssistant( - img=self.labels, + name="Pick characteristic points for segmentation.", + img=self.roi, background=self.background, - block=True, verbosity=self.verbosity, ) points = point_selection_assistant() + if points is not None and len(points) > 0: + points = np.fliplr(points) - labels = darsia.segment( + new_labels = darsia.segment( self.background, markers_method="supervised", edges_method="scharr", marker_points=points, mask=self.mask, region_size=2, + clean_up=False, ) if self.mask is not None: - # Use old values - old_labels = np.unique(self.labels.img[~self.mask]) - - # Determine holes in the old labels - missing_labels = np.setdiff1d( - np.unique(self.labels.img), np.unique(self.labels.img[~self.mask]) - ) - num_missing_labels = missing_labels.size - - # Fetch new values and map them on consecutive integers into old labels - unique_new_labels = np.unique(labels.img[self.mask]) - num_new_labels = unique_new_labels.size - new_labels = np.arange( - old_labels.max() + 1, - old_labels.max() + 1 + max(num_new_labels - num_missing_labels, 0), + # Determine a list of possible label ids - more than required. + # Consider all label ids marked by the mask, and additional ids extending + # the current label ids. Use a sufficient number of ids. + num_detected_labels = np.unique(new_labels.img[self.mask]).size + new_mapped_labels = np.concatenate( + ( + np.unique(self.labels.img[self.mask])[:num_detected_labels], + self.labels.img.max() + 1 + np.arange(max(0, num_detected_labels)), + ) ) - # Combine holes and new labels - unique_mapped_labels = np.concatenate((missing_labels, new_labels)) - - # Map new labels onto unique - for i, new_label in enumerate(unique_new_labels): - label_mask = np.isclose(labels.img, new_label) - labels.img[label_mask] = unique_mapped_labels[i] + # Assign new label ids + for i, new_label in enumerate(np.unique(new_labels.img[self.mask])): + label_mask = np.isclose(new_labels.img, new_label) + new_labels.img[label_mask] = new_mapped_labels[i] # Use old values in non-marked area - labels.img[~self.mask] = self.labels.img[~self.mask] + new_labels.img[~self.mask] = self.labels.img[~self.mask] + + return new_labels + + +class MonochromaticAssistant(darsia.BaseAssistant): + def __init__(self, img: darsia.Image, **kwargs) -> None: + self.name = "Monochromatic assistant" + """Name of the module.""" + self.input_img = img + """Input image.""" + self.img = img.to_monochromatic("gray") + """Monochromatic image.""" + + # Initialize base assistant with reference to monochromatic image + super().__init__(img=self.img, **kwargs) - return labels + self.finalized = False + """Flag indicating whether the assistant has been finalized.""" + + def _print_instructions(self) -> None: + print("") + print("Welcome to the monochromatic assistant.") + print("Please choose one of the following options:") + print(" - 'g': gray") + print(" - 'R': red") + print(" - 'G': green") + print(" - 'B': blue") + print(" - 'H': hue") + print(" - 'S': saturation") + print(" - 'V': value") + print(" - 'q': quit/abort\n") + + def _on_key_press(self, event) -> None: + if self.verbosity: + print(f"Current key: {event}") + + # Track events + if event.key == "g": + self.img = self.input_img.to_monochromatic("gray") + elif event.key == "R": + self.img = self.input_img.to_monochromatic("red") + elif event.key == "G": + self.img = self.input_img.to_monochromatic("green") + elif event.key == "B": + self.img = self.input_img.to_monochromatic("blue") + elif event.key == "H": + self.img = self.input_img.to_monochromatic("hue") + elif event.key == "S": + self.img = self.input_img.to_monochromatic("saturation") + elif event.key == "V": + self.img = self.input_img.to_monochromatic("value") + elif event.key == "q": + self.finalized = True + plt.close(self.fig) + + # Replot + if event.key in ["g", "R", "G", "B", "H", "S", "V"]: + self._plot_2d() + + def __call__(self) -> darsia.Image: + """Call the assistant.""" + + self._print_instructions() + plt.ion() + super().__call__() + plt.ioff() + return self.img class LabelsMaskSelectionAssistant: @@ -172,9 +260,9 @@ def __init__( def __call__(self) -> np.ndarray: # Identify points characterizing different regions to be merged point_selection_assistant = darsia.PointSelectionAssistant( + name="Mark labels.", img=self.labels, background=self.background, - block=True, verbosity=self.verbosity, ) points = point_selection_assistant() @@ -248,48 +336,67 @@ def __call__(self): class LabelsAssistant: def __init__( self, - labels: Optional[darsia.Image], + labels: Optional[darsia.Image] = None, background: Optional[darsia.Image] = None, **kwargs, ) -> None: self.name = "Labels assistant" """Name of the assistant.""" - self.labels = labels.copy() - """Input labels.""" - self.current_labels = labels - """Reference to current labels.""" self.background = background """Background image.""" - self.next_module = None - """Next module to be called.""" - self.finalized = False - """Flag indicating whether the assistant has been finalized.""" - self.verbosity = kwargs.get("verbosity", False) - - # Initialize empty labels - require background then - if self.labels is None: + self.monochromatic_background = None + """Monochromatic background image.""" + if isinstance(self.background, darsia.OpticalImage): + self.monochromatic_background = self.background.to_monochromatic("gray") + self.cache_background = None + """Cache for background image.""" + if labels is None: assert ( self.background is not None ), "Background image required to initialize empty labels." self.labels = darsia.Image( np.zeros_like(self.background.img, dtype=int), - **self.background.metadata, + **self.background.metadata(), ) + """Input labels.""" + else: + self.labels = labels.copy() + self.previous_labels = self.labels.copy() + """Reference to previous labels.""" + self.current_labels = self.labels.copy() + """Reference to current labels.""" + self.next_action = None + """Next module to be called.""" + self.finalized = False + """Flag indicating whether the assistant has been finalized.""" + self.verbosity = kwargs.get("verbosity", False) + """Verbosity flag.""" + self.first_call = True + """Flag indicating whether the assistant is called for the first time.""" def __call__(self): - """Call assistant.""" + """Call assistant. + + Always call the menu first. Then call the next module. Repeat until the + assistant is finalized. The instructions are printed only for the first call. + + Returns: + darsia.Image: current labels + + """ - # Initialize if not self.finalized: # Call menu self._call_menu() # Call next module - self._call_next_module() + self._call_next_action() # Repeat self.__call__() + self.first_call = False + # Return return self.current_labels @@ -301,38 +408,62 @@ def _call_menu(self) -> None: img=self.current_labels, background=self.background, verbosity=self.verbosity, + print_instructions=self.first_call, ) # Call menu module - self.next_module = self.menu() + self.next_action = self.menu() - def _call_next_module(self) -> None: + def _call_next_action(self) -> None: """Call next module.""" - if self.next_module == "segment": + # Fix previous labels if next action is modifying the labels + if self.next_action in ["segment", "merge", "refine"]: + self.previous_labels = self.current_labels.copy() + + # Apply action + if self.next_action == "segment": self._call_segment_module() - elif self.next_module == "pick": + elif self.next_action == "pick": self._call_pick_module() - elif self.next_module == "merge": + elif self.next_action == "merge": self._call_merge_module() - elif self.next_module == "refine": + elif self.next_action == "refine": self._call_refine_module() - elif self.next_module == "reset": + elif self.next_action == "reset": self.current_labels = self.labels.copy() - elif self.next_module == "quit": + elif self.next_action == "toggle_background": + self._toggle_background() + elif self.next_action == "monochromatic": + self._call_monochromatic_module() + elif self.next_action == "undo": + self.current_labels = self.previous_labels.copy() + elif self.next_action == "info": + self._print_info() + elif self.next_action == "quit": self.finalized = True + elif self.next_action is None: + pass else: - raise ValueError(f"Next module {self.next_module} not recognized.") + raise ValueError(f"Next module {self.next_action} not recognized.") + + if self.verbosity: + print(f"Next module: {self.next_action}") def _call_segment_module(self) -> None: """Call segment module.""" - segment_module = darsia.LabelsSegmentAssistant( - # labels=self.current_labels, - background=self.background, + if self.background is None: + assert self.cache_background is not None, "No background image available." + self.background = self.cache_background.copy() + self.cache_background = None + + segment_assistant = darsia.LabelsSegmentAssistant( + labels=None, + background=self.monochromatic_background, verbosity=self.verbosity, ) - self.current_labels = segment_module() + self.current_labels = segment_assistant() def _call_refine_module(self) -> None: """Call refine module.""" @@ -344,9 +475,14 @@ def _call_refine_module(self) -> None: ) mask = mask_selection_assistant() + if self.background is None: + assert self.cache_background is not None, "No background image available." + self.background = self.cache_background.copy() + self.cache_background = None + segment_assistant = darsia.LabelsSegmentAssistant( labels=self.current_labels, - background=self.background, + background=self.monochromatic_background, mask=mask, verbosity=self.verbosity, ) @@ -370,3 +506,35 @@ def _call_merge_module(self) -> None: background=self.background, ) self.current_labels = merge_assistant() + + def _print_info(self) -> None: + """Print out information about the assistant.""" + + print("The current labels:") + print(np.unique(self.current_labels.img).tolist()) + + # Plot current labels + plt.figure("Current labels") + plt.imshow(self.current_labels.img) + plt.show() + + def _toggle_background(self) -> None: + """Toggle background.""" + + if self.background is None and self.cache_background is None: + warn("No background image available.") + else: + if self.background is None: + self.background = self.cache_background.copy() + self.cache_background = None + else: + self.cache_background = self.background.copy() + self.background = None + + def _call_monochromatic_module(self) -> None: + if isinstance(self.background, darsia.OpticalImage): + monochromatic_assistant = darsia.MonochromaticAssistant( + img=self.background, + verbosity=self.verbosity, + ) + self.monochromatic_background = monochromatic_assistant()