53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
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()
|