134 lines
5.7 KiB
Python
134 lines
5.7 KiB
Python
"""Small, locally trainable colony-candidate detector for the PetriCount demo.
|
|
|
|
It deliberately uses only the annotated AGAR images that ship with this project.
|
|
The model ranks multiscale bright/dark blob candidates; it is a baseline for UI
|
|
integration, not a validated microbiology measurement method.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from scipy.ndimage import gaussian_filter, maximum_filter
|
|
|
|
MAX_SIDE = 720
|
|
|
|
|
|
@dataclass
|
|
class PreparedImage:
|
|
rgb: np.ndarray
|
|
gray: np.ndarray
|
|
scale: float
|
|
original_size: tuple[int, int]
|
|
|
|
|
|
def prepare_image(path_or_file) -> PreparedImage:
|
|
image = Image.open(path_or_file).convert("RGB")
|
|
original_size = image.size
|
|
scale = min(1.0, MAX_SIDE / max(image.size))
|
|
if scale < 1:
|
|
image = image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS)
|
|
rgb = np.asarray(image, dtype=np.float32) / 255.0
|
|
gray = rgb[..., 0] * 0.2126 + rgb[..., 1] * 0.7152 + rgb[..., 2] * 0.0722
|
|
return PreparedImage(rgb, gray, scale, original_size)
|
|
|
|
|
|
def _near_edge(x: int, y: int, width: int, height: int, margin: int = 13) -> bool:
|
|
return x < margin or y < margin or x >= width - margin or y >= height - margin
|
|
|
|
|
|
def candidate_points(prepared: PreparedImage) -> list[tuple[int, int]]:
|
|
"""Find potential round colonies as local extrema at several blob scales."""
|
|
gray = prepared.gray
|
|
height, width = gray.shape
|
|
candidates: list[tuple[float, int, int]] = []
|
|
# Colonies can be lighter or darker than the surrounding agar.
|
|
for sigma in (1.5, 2.5, 4.0, 6.0):
|
|
local = gaussian_filter(gray, sigma) - gaussian_filter(gray, sigma * 3.2)
|
|
radius = max(4, round(sigma * 2.5))
|
|
for response in (local, -local):
|
|
maxima = response == maximum_filter(response, size=radius * 2 + 1)
|
|
threshold = np.percentile(response, 97.2)
|
|
ys, xs = np.where(maxima & (response >= threshold))
|
|
for y, x in zip(ys, xs):
|
|
if _near_edge(int(x), int(y), width, height):
|
|
continue
|
|
# Suppress the metal/background outside the central dish region.
|
|
if np.hypot(x - width / 2, y - height / 2) > min(width, height) * 0.48:
|
|
continue
|
|
candidates.append((float(response[y, x]), int(x), int(y)))
|
|
candidates.sort(reverse=True)
|
|
selected: list[tuple[int, int]] = []
|
|
for _, x, y in candidates:
|
|
if all((x - px) ** 2 + (y - py) ** 2 > 13 ** 2 for px, py in selected):
|
|
selected.append((x, y))
|
|
if len(selected) >= 450:
|
|
break
|
|
# Diffuse colonies often have no sharp local maximum. Add a coarse candidate
|
|
# grid over visually changing parts of the dish and let the learned ranker
|
|
# decide whether each patch is a colony or background.
|
|
local_change = np.abs(gaussian_filter(gray, 2) - gaussian_filter(gray, 15))
|
|
change_threshold = np.percentile(local_change, 58)
|
|
for y in range(14, height - 14, 16):
|
|
for x in range(14, width - 14, 16):
|
|
if np.hypot(x - width / 2, y - height / 2) > min(width, height) * 0.46:
|
|
continue
|
|
if local_change[y, x] < change_threshold:
|
|
continue
|
|
if all((x - px) ** 2 + (y - py) ** 2 > 8 ** 2 for px, py in selected):
|
|
selected.append((x, y))
|
|
return selected
|
|
|
|
|
|
def _patch(array: np.ndarray, x: int, y: int, radius: int) -> np.ndarray:
|
|
return array[max(0, y - radius):y + radius + 1, max(0, x - radius):x + radius + 1]
|
|
|
|
|
|
def features_for_points(prepared: PreparedImage, points: Iterable[tuple[int, int]]) -> np.ndarray:
|
|
rows = []
|
|
smooth_small = gaussian_filter(prepared.gray, 2)
|
|
smooth_large = gaussian_filter(prepared.gray, 12)
|
|
contrast = smooth_small - smooth_large
|
|
gradient_y, gradient_x = np.gradient(prepared.gray)
|
|
gradient = np.hypot(gradient_x, gradient_y)
|
|
for x, y in points:
|
|
inner_gray = _patch(prepared.gray, x, y, 4)
|
|
outer_gray = _patch(prepared.gray, x, y, 12)
|
|
inner_rgb = _patch(prepared.rgb, x, y, 4).reshape(-1, 3)
|
|
outer_rgb = _patch(prepared.rgb, x, y, 12)
|
|
inner_gradient = _patch(gradient, x, y, 8)
|
|
rows.append([
|
|
*inner_rgb.mean(axis=0), *inner_rgb.std(axis=0),
|
|
float(inner_gray.mean()), float(inner_gray.std()),
|
|
float(outer_gray.mean()), float(outer_gray.std()),
|
|
float(contrast[y, x]), float(abs(contrast[y, x])),
|
|
float(inner_gradient.mean()), float(inner_gradient.std()),
|
|
x / prepared.gray.shape[1], y / prepared.gray.shape[0],
|
|
])
|
|
return np.asarray(rows, dtype=np.float32)
|
|
|
|
|
|
def label_centers(labels: list[dict], scale: float) -> list[tuple[int, int, float]]:
|
|
return [((item["x"] + item["width"] / 2) * scale, (item["y"] + item["height"] / 2) * scale,
|
|
max(item["width"], item["height"]) * scale / 2) for item in labels]
|
|
|
|
|
|
def label_candidates(points: list[tuple[int, int]], labels: list[dict], scale: float) -> np.ndarray:
|
|
centers = label_centers(labels, scale)
|
|
labels_for_points = []
|
|
for x, y in points:
|
|
labels_for_points.append(any((x - cx) ** 2 + (y - cy) ** 2 <= max(7, radius * .78) ** 2 for cx, cy, radius in centers))
|
|
return np.asarray(labels_for_points, dtype=np.int8)
|
|
|
|
|
|
def add_annotation_centers(points: list[tuple[int, int]], labels: list[dict], scale: float) -> list[tuple[int, int]]:
|
|
result = list(points)
|
|
for center_x, center_y, _ in label_centers(labels, scale):
|
|
point = (round(center_x), round(center_y))
|
|
if all((point[0] - px) ** 2 + (point[1] - py) ** 2 > 5 ** 2 for px, py in result):
|
|
result.append(point)
|
|
return result
|