utils by gemini

utils.py

@@ -1,163 +1,160 @@
-"""
-Collection of various utils 
-"""
-
 import numpy as np
+from PIL import Image, ImageDraw
+import logging
 
-import imageio.v3 as iio
-from PIL import Image
-# we may have very large images (e.g. panoramic SEM images), allow to read them w/o warnings
-Image.MAX_IMAGE_PIXELS = 933120000
-
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-from matplotlib.lines import Line2D
-
-
-import math
-
-
-###
-### load SEM images
-### 
-def load_image(filename : str) -> np.ndarray :
-    """Load an SEM image 
+def prepare_classifier_input(image, centroids, window_size):
+    """
+    Extracts image patches around centroids and prepares them as input for Keras models.
 
     Args:
-        filename (str): full path and name of the image file to be loaded
+        image: The input SEM image (PIL Image or NumPy array).
+        centroids (list): List of (x,y) coordinates of damage site centroids.
+        window_size (list): [height, width] of the square window to extract around each centroid.
 
     Returns:
-        np.ndarray: file as numpy ndarray
+        np.ndarray: A batch of image patches, ready for model prediction.
     """
-    image =  iio.imread(filename,mode='F')
-
-    return image
-
-
+    logging.info(f"prepare_classifier_input: Input image type: {type(image)}")
+
+    # Convert PIL Image to NumPy array if necessary
+    if isinstance(image, Image.Image):
+        # Convert to RGB first to ensure 3 channels for consistent model input
+        image_array = np.array(image.convert('RGB'))
+        logging.info("prepare_classifier_input: Converted PIL Image to RGB NumPy array.")
+    elif isinstance(image, np.ndarray):
+        # Ensure it's a 3-channel array for consistency if it's already NumPy
+        if image.ndim == 2: # Grayscale NumPy array
+            image_array = np.stack([image, image, image], axis=-1) # Convert to 3 channels
+            logging.info("prepare_classifier_input: Converted grayscale NumPy array to 3-channel.")
+        elif image.ndim == 3 and image.shape[2] == 4: # RGBA NumPy array
+            image_array = image[:, :, :3] # Drop alpha channel
+            logging.info("prepare_classifier_input: Dropped alpha channel from RGBA NumPy array.")
+        else: # Already RGB or similar 3-channel NumPy array
+            image_array = image
+            logging.info("prepare_classifier_input: Image is already a suitable NumPy array.")
+    else:
+        logging.error("prepare_classifier_input: Unsupported image format received. Expected PIL Image or NumPy array.")
+        raise ValueError("Unsupported image format for classifier input.")
+
+    if not centroids:
+        logging.warning("No centroids provided for prepare_classifier_input. Returning empty array.")
+        return np.empty((0, window_size[0], window_size[1], image_array.shape[2]), dtype=np.float32)
+
+    patches = []
+    img_height, img_width, _ = image_array.shape # Get dimensions from the now-guaranteed NumPy array
+    half_window_h, half_window_w = window_size[0] // 2, window_size[1] // 2
+
+    for c_y, c_x in centroids: # Centroids are (y, x) from clustering
+        # Ensure coordinates are integers
+        c_y, c_x = int(round(c_y)), int(round(c_x))
+
+        # Calculate bounding box for the patch
+        # Handle boundary conditions by clamping coordinates
+        y1 = max(0, c_y - half_window_h)
+        y2 = min(img_height, c_y + half_window_h)
+        x1 = max(0, c_x - half_window_w)
+        x2 = min(img_width, c_x + half_window_w)
+
+        # Extract patch
+        patch = image_array[y1:y2, x1:x2, :]
+
+        # Pad if the patch is smaller than window_size (due to boundary clamping)
+        if patch.shape[0] != window_size[0] or patch.shape[1] != window_size[1]:
+            padded_patch = np.zeros((window_size[0], window_size[1], image_array.shape[2]), dtype=patch.dtype)
+            padded_patch[0:patch.shape[0], 0:patch.shape[1], :] = patch
+            patch = padded_patch
+
+        patches.append(patch)
+    
+    # Normalize pixel values if your model expects it (e.g., to 0-1)
+    # This is a common step, adjust if your model's training pre-processing was different
+    # Assuming images are 0-255, converting to float 0-1
+    return np.array(patches, dtype=np.float32) / 255.0
 
-###
-### show SEM image with boxes in various colours around each damage site
-###
-def show_boxes(image : np.ndarray, damage_sites : dict, box_size = [250,250],
-               save_image = False, image_path : str = None) :
-    """_summary_
 
-    Args:
-        image (np.ndarray): SEM image to be shown
-        damage_sites (dict): python dictionary using the coordinates as key (x,y), and the label as value
-        box_size (list, optional): size of the rectangle drawn around each centroid. Defaults to [250,250].
-        save_image (bool, optional): save the image with the boxes or not. Defaults to False.
-        image_path (str, optional) : Full path and name of the output file to be saved
+def show_boxes(image, damage_sites, save_image=False, image_path="output_image.png"):
     """
+    Draws bounding boxes or markers on the image based on the classified damage sites.
 
-    _, ax = plt.subplots(1)
-    ax.imshow(image, cmap='gray')  # show image on correct axis
-    ax.set_xticks([])
-    ax.set_yticks([])
-
-    for key, label in damage_sites.items():
-        position = [key[0], key[1]]
-        edgecolor = {
-            'Inclusion': 'b',
-            'Interface': 'g',
-            'Martensite': 'r',
-            'Notch': 'y',
-            'Shadowing': 'm'
-        }.get(label, 'k')  # default: black
-
-        rect = patches.Rectangle((position[1] - box_size[1] / 2., position[0] - box_size[0] / 2),
-                                 box_size[1], box_size[0],
-                                 linewidth=1, edgecolor=edgecolor, facecolor='none')
-        ax.add_patch(rect)
-
-    legend_elements = [
-        Line2D([0], [0], color='b', lw=4, label='Inclusion'),
-        Line2D([0], [0], color='g', lw=4, label='Interface'),
-        Line2D([0], [0], color='r', lw=4, label='Martensite'),
-        Line2D([0], [0], color='y', lw=4, label='Notch'),
-        Line2D([0], [0], color='m', lw=4, label='Shadow'),
-        Line2D([0], [0], color='k', lw=4, label='Not Classified')
-    ]
-    ax.legend(handles=legend_elements, bbox_to_anchor=(1.04, 1), loc="upper left")
-
-    fig = ax.figure
-    fig.tight_layout(pad=0)
-
-    if save_image and image_path:
-        fig.savefig(image_path, dpi=1200, bbox_inches='tight')
-
-    canvas = fig.canvas
-    canvas.draw()
-
-    data = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8).reshape(
-        canvas.get_width_height()[::-1] + (4,))
-    data = data[:, :, :3]  # RGB only
-
-    plt.close(fig)
-
-    return data 
-
+    Args:
+        image: The input SEM image (PIL Image or NumPy array).
+        damage_sites (dict): Dictionary with (x,y) coordinates as keys and classification labels as values.
+        save_image (bool, optional): Whether to save the image to disk. Defaults to False.
+        image_path (str, optional): Path to save the image. Defaults to "output_image.png".
 
-###
-### cut out small images from panorama, append colour information
-###
-def prepare_classifier_input(panorama: np.ndarray, centroids: list, window_size=[250, 250]) -> list:
+    Returns:
+        PIL.Image.Image: The image with drawn boxes/markers.
     """
-    Extracts square image patches centered at each given centroid from a grayscale panoramic SEM image.
+    logging.info(f"show_boxes: Input image type: {type(image)}")
+
+    if image is None:
+        logging.warning("show_boxes received no image. Returning a blank image.")
+        img = Image.new('RGB', (500, 500), color = 'black')
+    else:
+        # Ensure image is a PIL Image for drawing operations
+        if isinstance(image, np.ndarray):
+            # Convert NumPy array to PIL Image. Assuming input is 0-255.
+            if image.dtype == np.float32 and np.max(image) <= 1.0: # If normalized 0-1 float
+                image_for_pil = (image * 255).astype(np.uint8)
+            else: # Assume 0-255 uint8
+                image_for_pil = image.astype(np.uint8)
+
+            if image_for_pil.ndim == 2: # Grayscale numpy
+                img = Image.fromarray(image_for_pil, mode='L').convert("RGB")
+            elif image_for_pil.ndim == 3 and image_for_pil.shape[2] in [3,4]: # RGB or RGBA
+                 img = Image.fromarray(image_for_pil).convert("RGB")
+            else:
+                 logging.error("Unsupported numpy image format for show_boxes.")
+                 img = Image.new('RGB', (500, 500), color = 'black') # Fallback
+        else: # Assume it's already a PIL Image
+            img = image.copy().convert("RGB") # Use a copy to avoid modifying original
+
+    draw = ImageDraw.Draw(img)
     
-    Each extracted patch is resized to the specified window size and converted into a 3-channel (RGB-like)
-    normalized image suitable for use with classification neural networks that expect color input.
-
-    Parameters
-    ----------
-    panorama : np.ndarray
-        Input SEM image. Should be a 2D array (H, W) or a 3D array (H, W, 1) representing grayscale data.
+    # Define some colors for drawing boxes
+    colors = {
+        "Inclusion": "red",
+        "Martensite": "blue",
+        "Interface": "green",
+        "Notch": "yellow",
+        "Shadowing": "purple",
+        "Not Classified": "gray", # Should ideally not appear on final image
+        "Unknown": "white"
+    }
     
-    centroids : list of [int, int]
-        List of (y, x) coordinates marking the centers of regions of interest. These are typically damage sites
-        identified in preprocessing (e.g., clustering).
-
-    window_size : list of int, optional
-        Size [height, width] of each extracted image patch. Defaults to [250, 250].
-
-    Returns
-    -------
-    list of np.ndarray
-        List of extracted and normalized 3-channel image patches, each with shape (height, width, 3). Only
-        centroids that allow full window extraction within image bounds are used.
-    """
-    if panorama.ndim == 2:
-        panorama = np.expand_dims(panorama, axis=-1)  # (H, W, 1)
-
-    H, W, _ = panorama.shape
-    win_h, win_w = window_size
-    images = []
-
-    for (cy, cx) in centroids:
-        x1 = int(cx - win_w / 2)
-        y1 = int(cy - win_h / 2)
-        x2 = x1 + win_w
-        y2 = y1 + win_h
-
-        # Skip if patch would go out of bounds
-        if x1 < 0 or y1 < 0 or x2 > W or y2 > H:
-            continue
-
-        # Extract and normalize patch
-        patch = panorama[y1:y2, x1:x2, 0].astype(np.float32)
-        patch = patch * 2. / 255. - 1.
-
-        # Replicate grayscale channel to simulate RGB
-        patch_color = np.repeat(patch[:, :, np.newaxis], 3, axis=2)
-        images.append(patch_color)
-
-    return images
-
+    for (x, y), label in damage_sites.items():
+        # Centroid coordinates from clustering (y,x) might be float
+        center_x = int(round(y)) # Note: (y,x) from clustering means y is row (height), x is column (width)
+        center_y = int(round(x)) # PIL expects (x, y) for drawing, so swap
 
+        box_size = 10 # Smaller box for clarity
+        
+        # Calculate box corners, clamping to image boundaries
+        x1 = max(0, center_x - box_size)
+        y1 = max(0, center_y - box_size)
+        x2 = min(img.width, center_x + box_size)
+        y2 = min(img.height, center_y + box_size)
 
+        fill_color = colors.get(label, "white")
+        outline_color = "black"
 
+        draw.rectangle([x1, y1, x2, y2], fill=fill_color, outline=outline_color, width=2)
+        
+        # Draw text label slightly offset
+        text_offset_x = 5
+        text_offset_y = -15
+        try:
+            draw.text((x1 + text_offset_x, y1 + text_offset_y), label, fill=outline_color)
+        except Exception as e:
+            logging.warning(f"Could not draw text label '{label}': {e}")
 
 
-        
+    if save_image and image_path:
+        try:
+            img.save(image_path)
+            logging.info(f"Image saved to {image_path}")
+        except Exception as e:
+            logging.error(f"Could not save image to {image_path}: {e}")
     
+    return img