Initial Bioportal Lab prototype

train_petri_model.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import joblib
+import numpy as np
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.metrics import precision_recall_fscore_support
+
+from petri_cv import add_annotation_centers, candidate_points, features_for_points, label_candidates, prepare_image
+
+DATASET = Path("AGAR_representative")
+MODEL_PATH = Path("models/petri_candidate_classifier.joblib")
+
+
+def main() -> None:
+    features, targets = [], []
+    annotations = sorted(DATASET.rglob("*.json"))
+    for annotation_path in annotations:
+        data = json.loads(annotation_path.read_text())
+        image_path = annotation_path.with_suffix(".jpg")
+        if not image_path.exists():
+            continue
+        image = prepare_image(image_path)
+        points = add_annotation_centers(candidate_points(image), data["labels"], image.scale)
+        labels = label_candidates(points, data["labels"], image.scale)
+        # keep a useful but not overwhelming background-to-colony ratio
+        positive = np.where(labels == 1)[0]
+        negative = np.where(labels == 0)[0]
+        rng = np.random.default_rng(abs(hash(annotation_path.stem)) % (2**32))
+        negative = rng.choice(negative, size=min(len(negative), max(50, len(positive) * 5)), replace=False)
+        selected = np.concatenate([positive, negative])
+        matrix = features_for_points(image, [points[index] for index in selected])
+        features.append(matrix)
+        targets.append(labels[selected])
+        print(f"{image_path.name}: {len(positive)} positives, {len(negative)} negatives")
+    x = np.vstack(features)
+    y = np.concatenate(targets)
+    model = RandomForestClassifier(
+        n_estimators=280, max_depth=14, min_samples_leaf=2, class_weight="balanced_subsample", n_jobs=-1, random_state=42
+    )
+    model.fit(x, y)
+    predicted = model.predict(x)
+    precision, recall, f1, _ = precision_recall_fscore_support(y, predicted, average="binary", zero_division=0)
+    MODEL_PATH.parent.mkdir(exist_ok=True)
+    joblib.dump({"model": model, "feature_version": 1}, MODEL_PATH)
+    print(f"Saved {MODEL_PATH}; training candidate precision={precision:.3f}, recall={recall:.3f}, f1={f1:.3f}; rows={len(y)}")
+
+
+if __name__ == "__main__":
+    main()