diff --git a/01_Scan_TopCamera_Rectangle.js b/01_Scan_TopCamera_Rectangle.js index 5945915..92ac87a 100644 --- a/01_Scan_TopCamera_Rectangle.js +++ b/01_Scan_TopCamera_Rectangle.js @@ -1105,25 +1105,65 @@ with (imports) { var stdoutLog = new File(controlDir, 'segmentation.out.log'); var stderrLog = new File(controlDir, 'segmentation.err.log'); var detectorMode = new File(controlDir, 'bug_detector.flag').exists() ? 'bug' : 'resistor'; + var classifierModel = new File(projectDir, 'models/insect_debris_classifier.pt'); try { - var builder = new Packages.java.lang.ProcessBuilder( - python, - segmentScript, - scanDir.getAbsolutePath(), - '--detector', - detectorMode, - '--watch' - ); + var builder; + if (detectorMode === 'bug') { + if (!classifierModel.exists()) { + throw new Error( + 'Bug detector classifier model not found: ' + + classifierModel.getAbsolutePath() + ); + } + + builder = new Packages.java.lang.ProcessBuilder( + python, + segmentScript, + scanDir.getAbsolutePath(), + '--detector', + detectorMode, + '--watch', + '--classifier-mode', + 'filter', + '--classifier-model', + classifierModel.getAbsolutePath(), + '--classifier-architecture', + 'efficientnet_b0', + '--classifier-threshold', + '0.50' + ); + } + else { + builder = new Packages.java.lang.ProcessBuilder( + python, + segmentScript, + scanDir.getAbsolutePath(), + '--detector', + detectorMode, + '--watch' + ); + } + builder.directory(projectDir); builder.redirectOutput(stdoutLog); builder.redirectError(stderrLog); builder.start(); - print('Launched ' + detectorMode + ' segmentation for: ' + scanDir.getAbsolutePath()); + + if (detectorMode === 'bug') { + print( + 'Launched bug segmentation with insect/debris classifier: ' + + classifierModel.getAbsolutePath() + ); + } + else { + print('Launched resistor segmentation for: ' + scanDir.getAbsolutePath()); + } } catch (error) { print('Failed to launch segmentation: ' + error); print('See: ' + stderrLog.getAbsolutePath()); + throw error; } } @@ -3198,14 +3238,34 @@ with (imports) { function waitForSegmentation(scanDir, timeoutMs) { var completeFile = new File(scanDir, 'segmentation_complete.json'); + var failedFile = new File(scanDir, 'segmentation_failed.json'); var start = new Date().getTime(); + while (!completeFile.exists()) { + if (failedFile.exists()) { + var failureMessage = 'Segmentation failed. See ' + failedFile.getAbsolutePath(); + try { + var failureRecord = JSON.parse(readText(failedFile)); + if (failureRecord.error) { + failureMessage += ': ' + failureRecord.error; + } + } + catch (failureReadError) { + print('Could not read segmentation failure details: ' + failureReadError); + } + + print(failureMessage); + return false; + } + if ((new Date().getTime() - start) > timeoutMs) { print('Timed out waiting for segmentation: ' + completeFile.getAbsolutePath()); return false; } + Packages.java.lang.Thread.sleep(500); } + return true; } diff --git a/02_Segment_Scan_Objects.py b/02_Segment_Scan_Objects.py index 459adac..baaa08a 100644 --- a/02_Segment_Scan_Objects.py +++ b/02_Segment_Scan_Objects.py @@ -58,6 +58,13 @@ raise SystemExit(2) from exc +try: + from insect_debris_classifier import InsectDebrisClassifier +except ModuleNotFoundError: + InsectDebrisClassifier = None + + + MIN_AREA_PX = 80 MAX_AREA_FRACTION = 0.03 MAX_RECT_AREA_FRACTION = 0.02 @@ -1306,6 +1313,12 @@ def csv_fields() -> list[str]: "absolute_background_contrast", "angle_degrees", "score", + "classifier_enabled", + "classifier_class", + "insect_probability", + "debris_probability", + "classifier_threshold", + "classifier_would_pick", ] @@ -1345,6 +1358,8 @@ def process_frame( csv_writer: csv.DictWriter, all_records: list[dict[str, Any]], args: argparse.Namespace, + classifier: Any | None = None, + rejected_file: Any | None = None, ) -> int: image_path = scan_dir / frame["file_name"] image = cv2.imread(str(image_path)) @@ -1397,6 +1412,49 @@ def process_frame( context_file=f"objects/{context_name}", overlay_file=overlay_file, ) + + record["classifier_enabled"] = classifier is not None + record["classifier_class"] = "" + record["insect_probability"] = None + record["debris_probability"] = None + record["classifier_threshold"] = None + record["classifier_would_pick"] = True + + if classifier is not None: + # Intentionally do not catch inference errors here. A classifier + # failure must stop segmentation rather than silently passing an + # unclassified candidate to the picker. + classification = classifier.classify_bgr(crop) + record["classifier_class"] = classification["class"] + record["insect_probability"] = classification["insect_probability"] + record["debris_probability"] = classification["debris_probability"] + record["classifier_threshold"] = classification["threshold"] + record["classifier_would_pick"] = classification["would_pick"] + + print( + f"Candidate {object_index}: " + f"{classification['class']} " + f"(insect={classification['insect_probability']:.3f}, " + f"debris={classification['debris_probability']:.3f}, " + f"threshold={classification['threshold']:.2f})", + flush=True, + ) + + if ( + args.classifier_mode == "filter" + and classifier is not None + and not bool(record["classifier_would_pick"]) + ): + # Keep rejected candidates for audit/debugging, but never place + # them in objects.jsonl/all_records where they could affect + # deduplication or become pick targets. + if rejected_file is not None: + rejected_file.write(json.dumps(record, sort_keys=True) + "\n") + rejected_file.flush() + + object_index += 1 + continue + all_records.append(record) frame_records.append(record) objects_file.write(json.dumps(record, sort_keys=True) + "\n") @@ -1548,8 +1606,33 @@ def scan_is_done(scan_dir: Path, processed_frames: int) -> bool: def run(args: argparse.Namespace) -> None: scan_dir = args.scan_dir + + segmentation_complete_path = scan_dir / "segmentation_complete.json" + segmentation_failed_path = scan_dir / "segmentation_failed.json" + segmentation_complete_path.unlink(missing_ok=True) + segmentation_failed_path.unlink(missing_ok=True) + calibration = load_training_tray_calibration(args.training_tray_calibration) + classifier = None + if args.classifier_mode != "off": + if args.detector != "bug": + raise ValueError( + "The insect/debris classifier can only be enabled with --detector bug." + ) + if InsectDebrisClassifier is None: + raise RuntimeError( + "Classifier support requires insect_debris_classifier.py and " + "torch/torchvision/Pillow in the BugPicker Python environment." + ) + + classifier = InsectDebrisClassifier( + model_path=args.classifier_model, + architecture=args.classifier_architecture, + threshold=args.classifier_threshold, + device=args.classifier_device, + ) + objects_dir = scan_dir / "objects" overlays_dir = scan_dir / "overlays" objects_dir.mkdir(exist_ok=True) @@ -1557,6 +1640,7 @@ def run(args: argparse.Namespace) -> None: objects_path = scan_dir / "objects.jsonl" csv_path = scan_dir / "objects.csv" + classifier_rejected_path = scan_dir / "classifier_rejected.jsonl" object_index = 0 all_records: list[dict[str, Any]] = [] processed_keys: set[tuple[int, str]] = set() @@ -1572,7 +1656,7 @@ def run(args: argparse.Namespace) -> None: with objects_path.open("w", encoding="utf-8") as objects_file, csv_path.open( "w", encoding="utf-8", newline="" - ) as csv_file: + ) as csv_file, classifier_rejected_path.open("w", encoding="utf-8") as rejected_file: csv_writer = csv.DictWriter(csv_file, fieldnames=csv_fields()) csv_writer.writeheader() @@ -1596,6 +1680,8 @@ def run(args: argparse.Namespace) -> None: csv_writer, all_records, args, + classifier, + rejected_file, ) processed_keys.add(key) @@ -1610,13 +1696,18 @@ def run(args: argparse.Namespace) -> None: time.sleep(args.poll_interval) - if object_index == 0: + if not all_records: + message = ( + "No pickable insect targets remained after classification" + if classifier is not None and args.classifier_mode == "filter" + else f"No {args.detector} targets detected in this segmentation run" + ) write_detection_status( scan_dir, "none_found", frame_index=max((key[0] for key in processed_keys), default=None), preview_file=None, - message=f"No {args.detector} targets detected in this segmentation run", + message=message, ) unique_object_count = 0 duplicate_count = 0 @@ -1676,9 +1767,76 @@ def main() -> int: ) parser.add_argument("--watch", action="store_true", help="Process frames as they are appended to manifest.jsonl") parser.add_argument("--poll-interval", type=float, default=0.5) + parser.add_argument( + "--classifier-mode", + choices=("off", "annotate", "filter"), + default="off", + help=( + "off: do not run the classifier; " + "annotate: classify candidates but keep all; " + "filter: reject debris before it can become a pick target." + ), + ) + parser.add_argument( + "--classifier-model", + type=Path, + default=PROJECT_ROOT / "models" / "insect_debris_classifier.pt", + help="Path to the trained insect/debris .pt checkpoint.", + ) + parser.add_argument( + "--classifier-architecture", + choices=("efficientnet_b0",), + default="efficientnet_b0", + help="Architecture used to train the checkpoint.", + ) + parser.add_argument( + "--classifier-threshold", + type=float, + default=0.50, + help="Minimum insect probability required to keep a candidate.", + ) + parser.add_argument( + "--classifier-device", + default="auto", + help="Torch device: auto, cpu, cuda, mps, etc.", + ) args = parser.parse_args() - run(args) + try: + run(args) + except Exception as exc: + failure_path = args.scan_dir / "segmentation_failed.json" + try: + failure_payload = { + "status": "failed", + "scan_dir": str(args.scan_dir), + "detector": args.detector, + "classifier_mode": args.classifier_mode, + "error": f"{type(exc).__name__}: {exc}", + "updated_at": datetime.now().isoformat(), + } + failure_path.write_text( + json.dumps(failure_payload, indent=2) + "\n", + encoding="utf-8", + ) + write_detection_status( + args.scan_dir, + "failed", + detector=args.detector, + message=f"Segmentation failed: {type(exc).__name__}: {exc}", + ) + except Exception as status_exc: + print( + f"Also failed to write segmentation failure status: {status_exc}", + file=sys.stderr, + ) + + print( + f"Segmentation failed: {type(exc).__name__}: {exc}", + file=sys.stderr, + ) + raise + return 0 diff --git a/insect_debris_classifier.py b/insect_debris_classifier.py new file mode 100644 index 0000000..8a070f3 --- /dev/null +++ b/insect_debris_classifier.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Runtime insect-vs-debris classifier for BugPicker detector crops.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import cv2 +import torch +import torch.nn as nn +from PIL import Image +from torchvision import models, transforms + + +CLASS_NAMES = ["debris", "insect"] +IMAGE_SIZE = 224 +DROPOUT = 0.20 + + +def select_device(requested: str = "auto") -> torch.device: + if requested != "auto": + return torch.device(requested) + + 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 build_model( + architecture: str, + number_of_classes: int, +) -> nn.Module: + if architecture == "efficientnet_b0": + model = models.efficientnet_b0(weights=None) + input_features = model.classifier[1].in_features + model.classifier = nn.Sequential( + nn.Dropout(p=DROPOUT), + nn.Linear(input_features, number_of_classes), + ) + return model + + raise ValueError( + f"Unsupported classifier architecture {architecture!r}. " + "This production patch currently supports efficientnet_b0." + ) + + +def build_inference_transform(): + return transforms.Compose( + [ + transforms.Resize(int(IMAGE_SIZE * 1.14)), + transforms.CenterCrop(IMAGE_SIZE), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ), + ] + ) + + +class InsectDebrisClassifier: + """Load one checkpoint and classify OpenCV BGR crops.""" + + def __init__( + self, + model_path: Path, + architecture: str = "efficientnet_b0", + threshold: float = 0.50, + device: str = "auto", + ) -> None: + self.model_path = Path(model_path).resolve() + if not self.model_path.exists(): + raise FileNotFoundError( + f"Insect/debris classifier checkpoint not found: {self.model_path}" + ) + + self.device = select_device(device) + self.threshold = float(threshold) + self.transform = build_inference_transform() + + checkpoint: Any = torch.load( + self.model_path, + map_location=self.device, + weights_only=False, + ) + + if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: + state_dict = checkpoint["model_state_dict"] + checkpoint_class_names = checkpoint.get("class_names") + elif isinstance(checkpoint, dict): + # Support a raw state_dict checkpoint. + state_dict = checkpoint + checkpoint_class_names = None + else: + raise TypeError( + "Classifier checkpoint must be a state_dict or a dictionary " + "containing model_state_dict." + ) + + self.class_names = list( + checkpoint_class_names + if checkpoint_class_names + else CLASS_NAMES + ) + + if "debris" not in self.class_names: + raise ValueError( + f"Checkpoint class_names must contain 'debris'; found {self.class_names}" + ) + + insect_label = None + for candidate in ("insect", "specimen"): + if candidate in self.class_names: + insect_label = candidate + break + + if insect_label is None: + raise ValueError( + "Checkpoint class_names must contain 'insect' or 'specimen'; " + f"found {self.class_names}" + ) + + self.debris_index = self.class_names.index("debris") + self.insect_index = self.class_names.index(insect_label) + + self.model = build_model( + architecture=architecture, + number_of_classes=len(self.class_names), + ) + self.model.load_state_dict(state_dict) + self.model.to(self.device) + self.model.eval() + + print( + "Loaded insect/debris classifier " + f"model={self.model_path} " + f"architecture={architecture} " + f"device={self.device} " + f"threshold={self.threshold:.2f} " + f"classes={self.class_names}", + flush=True, + ) + + def classify_bgr(self, image_bgr) -> dict[str, Any]: + if image_bgr is None or image_bgr.size == 0: + raise ValueError("Classifier received an empty image crop") + + image_rgb = cv2.cvtColor( + image_bgr, + cv2.COLOR_BGR2RGB, + ) + image = Image.fromarray(image_rgb) + + tensor = self.transform(image).unsqueeze(0).to(self.device) + + with torch.inference_mode(): + logits = self.model(tensor) + probabilities = torch.softmax(logits, dim=1)[0].detach().cpu() + + insect_probability = float(probabilities[self.insect_index]) + debris_probability = float(probabilities[self.debris_index]) + would_pick = insect_probability >= self.threshold + + return { + "class": "insect" if would_pick else "debris", + "insect_probability": insect_probability, + "debris_probability": debris_probability, + "threshold": self.threshold, + "would_pick": would_pick, + } diff --git a/models/insect_debris_classifier.pt b/models/insect_debris_classifier.pt new file mode 100644 index 0000000..36fd507 Binary files /dev/null and b/models/insect_debris_classifier.pt differ diff --git a/requirements.txt b/requirements.txt index cb3738c..81a3e3f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,5 @@ numpy opencv-python +torch +torchvision +Pillow