v0.4.0 release

README.md

@@ -57,6 +57,9 @@ Open the Pick tab. To the left is a a list of tomograms in the project, clicking
 
 ![Pick tab](docs/images/pick-tab.png)
 
+#### 3.1 Template matching
+Once several picks exist for a kind, the Match group within the settings pane can be used to identify potential candidates. Click `Scan` to generate an average pick and search unvisited tomograms for matches, which will appear as dotted markers on tomograms and in the pick table with a blue status icon. Click a candidate to select it, then press either `Accept` (or hotkey `a`) to convert it into a pick, or `Reject` (or either hotkey `delete`/`backspace`) to hide it. Currently this search is not rotation aware, so best suits symmetric particles. Setting for template matching can be edited in the project tab within the settings tab.
+
 ### 4. Review
 
 Open the Review tab. Picks can be filtered by kind, tomogram, minimum SNR and whether they were flagged, with options to delete the selected pick or all currently shown. All picks can be queued for re-estimation here if required as well. Single-click a pick to view XY/XZ/YZ previews, or double-click to open that pick in the Pick tab.
@@ -76,5 +79,5 @@ project/
   config.toml        settings and kinds
   tomograms.star     merged RELION tomogram table
   tilt_series/       one star file per tomogram
-  out/               picks, estimates, averages, RELION & .csv exports
+  out/               picks, candidates, estimates, averages, RELION & .csv exports
 ```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pickmix"
-version = "0.3.2"
+version = "0.4.0"
 description = "Multi-kind particle picking GUI for cryo-ET sub-tomogram averaging workflows."
 readme = "README.md"
 authors = [

src/pickmix/candidates.py

@@ -0,0 +1,149 @@
+'''
+pickmix: candidate picks found by template matching, saved to out_dir/candidates.csv
+'''
+
+# Import external dependencies
+import csv, itertools, os, threading
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from pydantic import BaseModel
+from typing import Literal
+
+# Import internal pickmix objects
+from pickmix.project import Project
+
+# Hit: candidate centre in reconstruction voxels and its score
+@dataclass(frozen=True)
+class Hit:
+    x: float
+    y: float
+    z: float
+    score: float
+
+# Candidate: region matching a template
+@dataclass
+class Candidate:
+    id: int
+    tomogram: str
+    kind: str
+    x: float
+    y: float
+    z: float
+    score: float
+    status: Literal['new', 'rejected'] = 'new'
+
+# _Row: one candidates.csv row
+class _Row(BaseModel):
+    tomogram: str
+    kind: str
+    x: float
+    y: float
+    z: float
+    score: float
+    status: Literal['new', 'rejected'] = 'new'
+
+# CandidateStore: every candidate, saved after each change
+class CandidateStore:
+    # __init__: load candidates.csv if there is one
+    def __init__(self, project: Project) -> None:
+        self.project = project
+        self.path: Path = project.config.out_dir / 'candidates.csv'
+        # called from whichever thread changed the store
+        self.listeners: list[Callable[[], None]] = []
+        self._lock = threading.Lock()
+        self._ids = itertools.count(1)
+        self._items: dict[int, Candidate] = {}
+        if self.path.exists():
+            with open(self.path, newline='') as handle:
+                for row in map(_Row.model_validate, csv.DictReader(handle)):
+                    candidate = Candidate(next(self._ids), **row.model_dump())
+                    self._items[candidate.id] = candidate
+
+    # get: the candidate with this id
+    def get(self, candidate_id: int) -> Candidate:
+        with self._lock:
+            return self._items[candidate_id]
+
+    # candidates: the new candidates in a tomogram, optionally of one kind and at or above a score, best first
+    def candidates(self, tomogram: str, kind: str | None = None, *, min_score: float = 0.0) -> list[Candidate]:
+        with self._lock:
+            found = [c for c in self._items.values() if c.status == 'new' and c.tomogram == tomogram and (kind is None or c.kind == kind) and c.score >= min_score]
+        return sorted(found, key=lambda c: -c.score)
+
+    # count: number of new candidates of a kind across all tomograms, or in one tomogram
+    def count(self, kind: str, tomogram: str | None = None, *, min_score: float = 0.0) -> int:
+        with self._lock:
+            return sum(c.status == 'new' and c.kind == kind and (tomogram is None or c.tomogram == tomogram) and c.score >= min_score for c in self._items.values())
+
+    # replace: the new candidates of a kind in a tomogram become these hits, except any within reach voxels of a rejected one
+    def replace(self, tomogram: str, kind: str, hits: Iterable[Hit], *, reach: float) -> None:
+        with self._lock:
+            rejected = [c for c in self._items.values() if c.status == 'rejected' and c.tomogram == tomogram and c.kind == kind]
+            for stale in [c.id for c in self._items.values() if c.status == 'new' and c.tomogram == tomogram and c.kind == kind]:
+                del self._items[stale]
+            for h in hits:
+                if all((h.x - r.x) ** 2 + (h.y - r.y) ** 2 + (h.z - r.z) ** 2 > reach**2 for r in rejected):
+                    candidate = Candidate(next(self._ids), tomogram, kind, h.x, h.y, h.z, h.score)
+                    self._items[candidate.id] = candidate
+        self._changed()
+
+    # reject: mark candidates as not wanted
+    def reject(self, *candidate_ids: int) -> None:
+        with self._lock:
+            for candidate_id in candidate_ids:
+                self._items[candidate_id].status = 'rejected'
+        self._changed()
+
+    # remove: forget candidates, as when they become picks
+    def remove(self, *candidate_ids: int) -> None:
+        with self._lock:
+            for candidate_id in candidate_ids:
+                del self._items[candidate_id]
+        self._changed()
+
+    # clear: forget the candidates of a kind, or all, rejected ones included
+    def clear(self, kind: str | None = None) -> None:
+        with self._lock:
+            self._items = {i: c for i, c in self._items.items() if kind is not None and c.kind != kind}
+        self._changed()
+
+    # prune: forget every candidate for a tomogram
+    def prune(self, tomograms: Iterable[str]) -> None:
+        gone = set(tomograms)
+        with self._lock:
+            kept = {i: c for i, c in self._items.items() if c.tomogram not in gone}
+            changed = len(kept) != len(self._items)
+            self._items = kept
+        if changed:
+            self._changed()
+
+    # rename_kind: move every candidate of a kind to a new name
+    def rename_kind(self, old: str, new: str) -> None:
+        with self._lock:
+            for candidate in self._items.values():
+                if candidate.kind == old:
+                    candidate.kind = new
+        self._changed()
+
+    # remove_kind: forget every candidate of a kind
+    def remove_kind(self, kind: str) -> None:
+        self.clear(kind)
+
+    # save: write CSV atomically
+    def save(self) -> None:
+        with self._lock:
+            rows = [_Row(**{k: v for k, v in vars(c).items() if k != 'id'}).model_dump() for c in self._items.values()]
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = self.path.with_suffix('.csv.tmp')
+        with open(tmp, 'w', newline='') as handle:
+            writer = csv.DictWriter(handle, fieldnames=list(_Row.model_fields))
+            writer.writeheader()
+            writer.writerows(rows)
+        os.replace(tmp, self.path)
+
+    # _changed: save, then tell listeners
+    def _changed(self) -> None:
+        self.save()
+        for listener in self.listeners:
+            listener()

src/pickmix/config.py

@@ -51,6 +51,21 @@ class RingParams(BaseModel):
     # box edge as a multiple of the ring diameter when the kind takes its box from the estimator
     box_margin: float = Field(default=1.25, title='Box margin (× diameter)', description='Box edge as a multiple of the ring diameter, for kinds that take their box from the estimator.')
 
+# MatchParams: settings of the template matching scan; title and description feed the settings form
+class MatchParams(BaseModel):
+    # template is binned until its edge is about this many pixels, which sets the speed of the scan
+    template_px: int = Field(default=32, ge=8, title='Template edge (px)', description='The tomogram and template are binned until the template is about this many pixels across. Smaller is faster and coarser.')
+    # candidates scoring below this are not kept
+    min_score: float = Field(default=0.2, gt=0, lt=1, title='Minimum score', description='Candidates with a normalised cross-correlation below this are not kept.')
+    # nearest two candidates may be, as a multiple of the kind's marker diameter
+    separation: float = Field(default=0.75, gt=0, title='Separation (× marker diameter)', description='Closest two candidates, or a candidate and a pick, may be.')
+    # mask radius as a multiple of the marker radius
+    mask_margin: float = Field(default=1.2, gt=0, title='Mask radius (× marker radius)', description='Radius of the spherical mask applied to the template, as a multiple of the marker radius.')
+    # best candidates kept per tomogram
+    max_per_tomogram: int = Field(default=200, ge=1, title='Most per tomogram', description='Most candidates kept in one tomogram, best first.')
+    # threads the correlation may use
+    workers: int = Field(default=2, ge=1, title='Threads', description='Threads the correlation may use. More is faster but slows the rest of the program.')
+
 # _default_kinds: the single cage kind used when a config has no kinds table
 def _default_kinds() -> dict[str, Kind]:
     return {'cage': Kind(marker_a=600.0, estimator='ring')}
@@ -73,6 +88,7 @@ class Config(BaseModel):
     # most memory the slice cache may hold; scrolling is smooth where a whole tomogram fits
     cache_mb: int = Field(default=2048, ge=64, title='Slice cache (MB)', description='Most memory the viewer may use to keep slices of tomograms ready for scrolling.')
     ring: RingParams = Field(default_factory=RingParams)
+    match: MatchParams = Field(default_factory=MatchParams)
     kinds: dict[str, Kind] = Field(default_factory=_default_kinds)
     # kind chosen by each of the keys 1 to 9, '' for an unused key
     hotkeys: list[str] = Field(default_factory=list)

src/pickmix/gui/components/canvas.py

@@ -45,6 +45,9 @@ class Marker:
     fit: FitView | None = None
     # membrane normal projected on the slice, as a unit-length-or-shorter (x, y) vector
     normal: tuple[float, float] | None = None
+    # template matching candidate rather than a pick
+    candidate: bool = False
+    tip: str = ''
 
 # SliceCanvas: one slice plus markers; emits requests, never edits picks itself
 class SliceCanvas(QGraphicsView):
@@ -151,9 +154,10 @@ def set_markers(self, markers: list[Marker]) -> None:
             ring = QGraphicsEllipseItem(-m.radius_px, -m.radius_px, 2 * m.radius_px, 2 * m.radius_px)
             ring.setPos(m.x + 0.5, m.y + 0.5)
             ring.setZValue(1)
+            ring.setToolTip(m.tip)
             dot = QGraphicsEllipseItem(-2, -2, 4, 4, ring)
             dot.setFlag(QGraphicsEllipseItem.GraphicsItemFlag.ItemIgnoresTransformations)
-            dot.setBrush(QColor(m.colour))
+            dot.setBrush(Qt.BrushStyle.NoBrush if m.candidate else QColor(m.colour))
             dot.setPen(QPen(Qt.PenStyle.NoPen))
             if m.normal is not None:
                 reach = max(4 * m.radius_px, 16.0)
@@ -212,6 +216,9 @@ def mouseMoveEvent(self, event: QMouseEvent) -> None:
             self.verticalScrollBar().setValue(self.verticalScrollBar().value() - int(delta.y()))
         elif self._drag is not None:
             pick_id, start = self._drag
+            # don't drag a candidate
+            if self._markers[pick_id].candidate:
+                return
             pos = self.mapToScene(event.position().toPoint())
             delta = pos - start
             if self._drag_moved or math.hypot(delta.x(), delta.y()) * self.transform().m11() > _DRAG_PX:
@@ -275,6 +282,8 @@ def _restyle(self) -> None:
             pen.setWidth(3 if marker_id == self.selected_id else 2)
             if m.flagged:
                 pen.setStyle(Qt.PenStyle.DashLine)
+            elif m.candidate:
+                pen.setStyle(Qt.PenStyle.DotLine)
             item.setPen(pen)
         for arrow in self._arrows.values():
             arrow.setVisible(self._normals)

src/pickmix/gui/components/status.py

@@ -12,15 +12,16 @@
 from pickmix.gui.components.icons import tinted
 from pickmix.gui.style import ERROR, OK, RUNNING, WARN
 
-# Status: ok is a green tick, warn a yellow warning sign, error a red error sign, running a blue spinner
-type Status = Literal['ok', 'warn', 'error', 'running']
+# Status: ok is a green tick, warn a yellow warning sign, error a red error sign, running a blue spinner, match a blue image search for a template match candidate
+type Status = Literal['ok', 'warn', 'error', 'running', 'match']
 
 # STATUS_STYLE: icon stem, colour and the text used when the icon file is missing
 STATUS_STYLE: dict[Status, tuple[str, str, str]] = {
     'ok': ('check', OK, '✓'),
     'warn': ('warning', WARN, '!'),
     'error': ('error', ERROR, '✖'),
     'running': ('refresh', RUNNING, '…'),
+    'match': ('image_search', RUNNING, 'tm'),
 }
 
 # _SIZE: edge in pixels of a status icon

src/pickmix/gui/main_window.py

@@ -16,6 +16,7 @@
 from pickmix.gui.tabs.project_tab import ProjectTab
 from pickmix.gui.tabs.review_tab import ReviewTab
 from pickmix.loader import Loader
+from pickmix.matcher import MatchService
 from pickmix.model import Session
 from pickmix.project import Project
 
@@ -39,6 +40,7 @@ def __init__(self) -> None:
         self.session: Session | None = None
         self.loader: Loader | None = None
         self.service: EstimateService | None = None
+        self.matcher: MatchService | None = None
         self.tabs = QTabWidget()
         self.project_tab = ProjectTab()
         self.kinds_pane = self.project_tab.kinds_pane
@@ -70,6 +72,7 @@ def __init__(self) -> None:
         self.toast = install(self)
         for i in range(1, 10):
             QShortcut(QKeySequence(str(i)), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(lambda i=i: self._hotkey(i - 1))
+        QShortcut(QKeySequence('A'), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(self._accept_candidate)
         for key in ('Delete', 'Backspace'):
             QShortcut(QKeySequence(key), self, context=Qt.ShortcutContext.WindowShortcut).activated.connect(self._delete_pick)
         self.resize(*WELCOME_SIZE)
@@ -107,6 +110,11 @@ def _hotkey(self, i: int) -> None:
         if self._picking():
             self.pick_tab.pick_kind(i)
 
+    # _accept_candidate: A converts selected candidate to pick
+    def _accept_candidate(self) -> None:
+        if self._picking():
+            self.pick_tab.accept_selected()
+
     # _delete_pick: Delete or Backspace removes the selected pick, from wherever focus is
     def _delete_pick(self) -> None:
         if self._picking():
@@ -162,12 +170,13 @@ def _on_project(self, project: Project) -> None:
         if project.tomos:
             self.session, self.loader = Session(project), Loader(project)
             self.service = EstimateService(self.session, self.loader)
-            self.pick_tab.set_session(self.session, self.loader, self.service)
+            self.matcher = MatchService(self.session, self.service.store)
+            self.pick_tab.set_session(self.session, self.loader, self.service, self.matcher)
             self.review_tab.set_session(self.session, self.loader, self.service)
             self.export_tab.set_session(self.session, self.loader, self.service)
             self.kinds_pane.set_session(self.session, self.service)
         else:
-            self.session = self.loader = self.service = None
+            self.session = self.loader = self.service = self.matcher = None
             self.kinds_pane.set_session(None, None)
             self.export_tab.clear()
         self._set_enabled(session=self.session is not None)
@@ -176,9 +185,11 @@ def _on_project(self, project: Project) -> None:
     def _shutdown(self) -> None:
         if self.service is not None:
             self.service.shutdown()
+        if self.matcher is not None:
+            self.matcher.shutdown()
         if self.loader is not None:
             self.loader.shutdown()
-        self.session = self.loader = self.service = None
+        self.session = self.loader = self.service = self.matcher = None
 
     # _on_closed: no project; lock everything but the Project tab
     def _on_closed(self) -> None:

src/pickmix/gui/panes/kinds_pane.py

@@ -186,6 +186,7 @@ def _apply(self, name: str, kind: Kind, *, old: str | None = None) -> None:
                 self.service.store.rename_kind(old, name)
             if self.session is not None:
                 self.session.store.rename_kind(old, name)
+                self.session.candidates.rename_kind(old, name)
         self._saved()
 
     # _delete: remove a kind and its picks, asking first if it has any; refused with a reason if the config would not be valid
@@ -204,6 +205,7 @@ def _delete(self, name: str) -> None:
         self.project.update(**changes)
         if self.session is not None:
             self.session.store.remove_kind(name)
+            self.session.candidates.remove_kind(name)
         self._saved()
 
     # _confirm: ask a yes or no question