diff --git a/01_Scan_TopCamera_Rectangle.js b/01_Scan_TopCamera_Rectangle.js index 598e156..b2a8e6c 100644 --- a/01_Scan_TopCamera_Rectangle.js +++ b/01_Scan_TopCamera_Rectangle.js @@ -3576,6 +3576,50 @@ with (imports) { } } + function runPlateTaxonomyClassifier(plateContexts, controlDir) { + var taxonomyScript = new File(scriptsDir, '08_Classify_Plate_Taxonomy.py').getAbsolutePath(); + var allSucceeded = true; + + for (var taxonomyIndex = 0; taxonomyIndex < plateContexts.length; taxonomyIndex++) { + var taxonomyContext = plateContexts[taxonomyIndex]; + var taxonomyPlate = normalizePlateNumber(taxonomyContext.plateNumber); + var stdoutLog = new File(controlDir, 'taxonomy_classifier_' + taxonomyPlate + '.out.log'); + var stderrLog = new File(controlDir, 'taxonomy_classifier_' + taxonomyPlate + '.err.log'); + + try { + var builder = new Packages.java.lang.ProcessBuilder( + python, + taxonomyScript, + '--openpnp-root', + projectDir.getAbsolutePath(), + '--plate', + taxonomyPlate + ); + builder.directory(projectDir); + builder.redirectOutput(stdoutLog); + builder.redirectError(stderrLog); + var process = builder.start(); + var exitCode = process.waitFor(); + if (exitCode === 0) { + print('Taxonomy classification completed for plate ' + taxonomyPlate + '.'); + } + else { + allSucceeded = false; + print('Taxonomy classifier exited with code ' + exitCode + ' for plate ' + taxonomyPlate + + '. Plate data is preserved and the run will continue.'); + print('See: ' + stderrLog.getAbsolutePath()); + } + } + catch (error) { + allSucceeded = false; + print('Failed to run taxonomy classifier for plate ' + taxonomyPlate + ': ' + error); + print('See: ' + stderrLog.getAbsolutePath()); + } + } + + return allSucceeded; + } + function touchTargets(scanDir, pauseFile, stopFile, statusFile, scanId, totalFrames) { var touchHeadName = 'H1'; var touchNozzleName = 'N1'; @@ -7386,6 +7430,8 @@ with (imports) { && pickResult !== null && pickResult.attemptedWells > 0 && !touchDryRunFile.exists()) { + writeStatus(statusFile, 'classifying', scanId, totalFrames, totalFrames, 'Running BioCLIP taxonomy classification'); + runPlateTaxonomyClassifier(plateContexts, controlDir); var completedAuditMessages = []; for (var auditIndex = 0; auditIndex < plateContexts.length; auditIndex++) { var auditContext = plateContexts[auditIndex]; diff --git a/08_Classify_Plate_Taxonomy.py b/08_Classify_Plate_Taxonomy.py new file mode 100644 index 0000000..7d7c1fe --- /dev/null +++ b/08_Classify_Plate_Taxonomy.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""08: Classify completed BugPicker plate specimens from their saved multi-view images.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import os +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from taxonomy_classifier import TaxonomyClassifier + +IMAGE_KINDS = ("scan", "well", "bottom", "hires") +INVALID_IMAGE_CODES = {"", "nan", "none", "null", "", "blank"} + +AI_COLUMNS = [ + "AI Order", + "AI Order Confidence", + "AI Order Top 3", + "AI Species", + "AI Species Confidence", + "AI Species Top 3", + "AI Taxonomy Model", + "AI Taxonomy Fusion", + "AI Taxonomy Views", + "AI Taxonomy Status", + "AI Taxonomy Error", + "AI Taxonomy Timestamp", +] + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def normalize_plate(value: str) -> str: + plate = "".join(ch for ch in value.strip().upper() if ch.isalnum() or ch in "_-") + if not plate: + raise ValueError("Plate number must not be blank") + return plate + + +def normalize_image_code(value: Any) -> str: + """Return a usable Image Code, or an empty string for placeholder/blank cells.""" + if value is None: + return "" + image_code = str(value).strip() + if image_code.lower() in INVALID_IMAGE_CODES: + return "" + return image_code + + +def find_image_paths(image_dir: Path, image_code: str) -> dict[str, Path]: + result: dict[str, Path] = {} + for kind in IMAGE_KINDS: + path = image_dir / f"{image_code}_{kind}.png" + if path.exists(): + result[kind] = path + return result + + +def format_top3(items: list[dict[str, Any]]) -> str: + return " | ".join(f"{item['name']}:{float(item['confidence']):.6f}" for item in items) + + +def mock_prediction(image_code: str, views: list[str]) -> dict[str, Any]: + orders = ["Diptera", "Hemiptera", "Hymenoptera", "Coleoptera"] + digest = hashlib.sha256(image_code.encode("utf-8")).digest() + index = digest[0] % len(orders) + order = orders[index] + confidence = 0.80 + ((digest[1] % 16) / 100.0) + other = [value for value in orders if value != order][:2] + top3 = [ + {"name": order, "confidence": confidence}, + {"name": other[0], "confidence": (1.0 - confidence) * 0.65}, + {"name": other[1], "confidence": (1.0 - confidence) * 0.35}, + ] + return { + "model": "mock_taxonomy", + "kind": "mock", + "fusion": "mock_four_view", + "checkpoint": "", + "views_used": views, + "order": order, + "order_confidence": confidence, + "order_top3": top3, + "species": "mock species", + "species_confidence": confidence * 0.85, + "species_top3": [{"name": "mock species", "confidence": confidence * 0.85}], + } + + +def atomic_write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temp_name = tempfile.mkstemp(prefix=path.name + ".", suffix=".tmp", dir=path.parent) + os.close(fd) + temp_path = Path(temp_name) + try: + with temp_path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + temp_path.replace(path) + finally: + if temp_path.exists(): + temp_path.unlink() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--openpnp-root", type=Path, required=True) + parser.add_argument("--plate", required=True) + parser.add_argument("--force", action="store_true") + parser.add_argument("--mock", action="store_true", help="Exercise integration without loading BioCLIP/checkpoints") + args = parser.parse_args() + + root = args.openpnp_root.resolve() + plate = normalize_plate(args.plate) + csv_path = root / "Plate_spreadsheets" / f"{plate}.csv" + image_dir = root / "Plate_insect_images" / f"P-{plate}" + if not csv_path.exists(): + raise SystemExit(f"Plate CSV not found: {csv_path}") + if not image_dir.exists(): + raise SystemExit(f"Plate image directory not found: {image_dir}") + + with csv_path.open(encoding="utf-8", newline="") as handle: + reader = csv.DictReader(handle) + fieldnames = list(reader.fieldnames or []) + rows = [dict(row) for row in reader] + for column in AI_COLUMNS: + if column not in fieldnames: + fieldnames.append(column) + + classifier = None if args.mock else TaxonomyClassifier() + processed = 0 + skipped = 0 + errors = 0 + run_records: list[dict[str, Any]] = [] + + for row in rows: + image_code = normalize_image_code(row.get("Image Code")) + well = str(row.get("Well Number") or "").strip() + if not image_code: + skipped += 1 + continue + if not args.force and str(row.get("AI Taxonomy Status") or "").strip().lower() == "classified": + skipped += 1 + continue + + image_paths = find_image_paths(image_dir, image_code) + try: + if not image_paths: + raise FileNotFoundError(f"No plate specimen images found for {image_code}") + prediction = ( + mock_prediction(image_code, sorted(image_paths)) + if args.mock + else classifier.classify(image_paths) # type: ignore[union-attr] + ) + row["AI Order"] = str(prediction["order"]) + row["AI Order Confidence"] = f"{float(prediction['order_confidence']):.6f}" + row["AI Order Top 3"] = format_top3(prediction["order_top3"]) + row["AI Species"] = str(prediction.get("species", "")) + row["AI Species Confidence"] = f"{float(prediction.get('species_confidence', 0.0)):.6f}" + row["AI Species Top 3"] = format_top3(prediction.get("species_top3", [])) + row["AI Taxonomy Model"] = str(prediction["model"]) + row["AI Taxonomy Fusion"] = str(prediction["fusion"]) + row["AI Taxonomy Views"] = "+".join(prediction["views_used"]) + row["AI Taxonomy Status"] = "classified" + row["AI Taxonomy Error"] = "" + row["AI Taxonomy Timestamp"] = utc_now() + processed += 1 + run_records.append({"plate": plate, "well": well, "image_code": image_code, **prediction}) + print( + f"{image_code}: {prediction['order']} " + f"({float(prediction['order_confidence']):.2%}) " + f"views={'+'.join(prediction['views_used'])} model={prediction['model']}", + flush=True, + ) + except Exception as exc: + errors += 1 + row["AI Taxonomy Status"] = "error" + row["AI Taxonomy Error"] = str(exc) + row["AI Taxonomy Timestamp"] = utc_now() + run_records.append({ + "plate": plate, + "well": well, + "image_code": image_code, + "status": "error", + "error": str(exc), + }) + print(f"ERROR {image_code}: {exc}", flush=True) + + atomic_write_csv(csv_path, fieldnames, rows) + image_dir.mkdir(parents=True, exist_ok=True) + last_run = image_dir / "taxonomy_predictions_last_run.jsonl" + with last_run.open("w", encoding="utf-8") as handle: + for record in run_records: + handle.write(json.dumps(record, sort_keys=True) + "\n") + history = image_dir / "taxonomy_predictions_history.jsonl" + with history.open("a", encoding="utf-8") as handle: + for record in run_records: + history_record = dict(record) + history_record["recorded_at"] = utc_now() + handle.write(json.dumps(history_record, sort_keys=True) + "\n") + + summary = { + "plate": plate, + "csv_path": str(csv_path), + "image_dir": str(image_dir), + "mock": bool(args.mock), + "processed": processed, + "skipped": skipped, + "errors": errors, + "rows": len(rows), + } + (image_dir / "taxonomy_predictions_summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(summary, sort_keys=True), flush=True) + return 0 if errors == 0 else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/Data/taxonomy_classifier.pt b/Data/taxonomy_classifier.pt new file mode 100644 index 0000000..1befa25 Binary files /dev/null and b/Data/taxonomy_classifier.pt differ diff --git a/requirements.txt b/requirements.txt index 9db976b..720997e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ opencv-python pillow torch torchvision +open_clip_torch diff --git a/taxonomy_classifier.py b/taxonomy_classifier.py new file mode 100644 index 0000000..866e37e --- /dev/null +++ b/taxonomy_classifier.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""BioCLIP multi-view taxonomy inference using the deployed hierarchical model.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from PIL import Image + + +SCRIPT_DIR = Path(__file__).resolve().parent +BIOCLIP_MODEL = "hf-hub:imageomics/bioclip-2" +CHECKPOINT_PATH = SCRIPT_DIR / "Data" / "taxonomy_classifier.pt" + + +class HierarchicalClassifier(nn.Module): + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + num_species: int, + num_orders: int, + dropout: float, + ) -> None: + super().__init__() + self.shared = nn.Sequential( + nn.Linear(embedding_dim, hidden_dim), + nn.ReLU(), + ) + self.dropout = nn.Dropout(dropout) + self.species_head = nn.Linear(hidden_dim, num_species) + self.order_head = nn.Linear(hidden_dim, num_orders) + + def forward(self, embeddings: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + hidden = self.dropout(self.shared(embeddings)) + return self.species_head(hidden), self.order_head(hidden) + + +def select_device() -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda") + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + +def normalize_probabilities(probabilities: torch.Tensor) -> torch.Tensor: + total = probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12) + return probabilities / total + + +def top_predictions( + probabilities: torch.Tensor, + names: list[str], + k: int = 3, +) -> list[dict[str, Any]]: + probabilities = probabilities.detach().float().cpu().reshape(-1) + k = min(k, len(names)) + scores, indices = torch.topk(probabilities, k=k) + return [ + {"name": names[int(index)], "confidence": float(score)} + for score, index in zip(scores.tolist(), indices.tolist()) + ] + + +class TaxonomyClassifier: + """Run the deployed iNaturalist-trained hierarchical BioCLIP classifier.""" + + def __init__(self) -> None: + self.model_name = BIOCLIP_MODEL + self.device = select_device() + self.checkpoint_path = CHECKPOINT_PATH + + if not self.checkpoint_path.exists(): + raise FileNotFoundError( + f"Taxonomy checkpoint not found: {self.checkpoint_path}\n" + "Copy taxonomy_classifier.pt into the BugPicker Data directory." + ) + + checkpoint: dict[str, Any] = torch.load( + self.checkpoint_path, + map_location="cpu", + weights_only=False, + ) + + self.species_names = list(checkpoint["species_names"]) + self.order_names = list(checkpoint["all_orders"]) + + embedding_dim = int(checkpoint.get("embedding_dim", 768)) + hidden_dim = int(checkpoint.get("hidden_dim", 512)) + dropout = float(checkpoint.get("dropout", 0.30)) + + self.classifier = HierarchicalClassifier( + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + num_species=len(self.species_names), + num_orders=len(self.order_names), + dropout=dropout, + ) + self.classifier.load_state_dict(checkpoint["model_state_dict"]) + self.classifier = self.classifier.to(self.device).eval() + + try: + import open_clip + except ModuleNotFoundError as exc: + raise RuntimeError( + "open_clip is required for taxonomy inference. " + "Install open_clip_torch from requirements.txt." + ) from exc + + self.bioclip, _, self.preprocess = open_clip.create_model_and_transforms( + self.model_name + ) + self.bioclip = self.bioclip.to(self.device).eval() + for parameter in self.bioclip.parameters(): + parameter.requires_grad = False + + def _embed_paths( + self, + image_paths: dict[str, Path], + ) -> tuple[list[str], torch.Tensor]: + view_names = sorted(image_paths) + tensors = [] + + for view in view_names: + with Image.open(image_paths[view]) as image: + tensors.append(self.preprocess(image.convert("RGB"))) + + images = torch.stack(tensors).to(self.device) + + with torch.inference_mode(): + embeddings = self.bioclip.encode_image(images).float() + embeddings = F.normalize(embeddings, p=2, dim=-1) + + return view_names, embeddings + + def classify(self, image_paths: dict[str, Path]) -> dict[str, Any]: + if not image_paths: + raise ValueError("No specimen images were supplied") + + view_names, embeddings = self._embed_paths(image_paths) + + # Best-performing deployment path from the controlled BioKEA comparison: + # direct hierarchical Order head + probability averaging across views. + with torch.inference_mode(): + _, order_logits_by_view = self.classifier(embeddings) + order_probs_by_view = torch.softmax(order_logits_by_view, dim=-1) + order_probs = normalize_probabilities(order_probs_by_view.mean(dim=0)) + + # Keep a species prediction from the same hierarchical checkpoint. + # Species uses an averaged BioCLIP embedding and is informational only + # because BioKEA currently has Order-level ground truth. + averaged_embedding = F.normalize( + embeddings.mean(dim=0, keepdim=True), + p=2, + dim=-1, + ) + species_logits, _ = self.classifier(averaged_embedding) + species_probs = torch.softmax(species_logits, dim=-1).squeeze(0) + + order_top3 = top_predictions(order_probs, self.order_names, k=3) + species_top3 = top_predictions(species_probs, self.species_names, k=3) + + return { + "model": "taxonomy_classifier", + "kind": "hierarchical_direct_order", + "fusion": "order_probability_average", + "checkpoint": self.checkpoint_path.name, + "views_used": view_names, + "order": order_top3[0]["name"], + "order_confidence": order_top3[0]["confidence"], + "order_top3": order_top3, + "species": species_top3[0]["name"], + "species_confidence": species_top3[0]["confidence"], + "species_top3": species_top3, + }