Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import sys | |
| from pathlib import Path | |
| import numpy as np | |
| from sklearn.linear_model import LogisticRegression | |
| from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score | |
| from sklearn.pipeline import Pipeline | |
| from sklearn.preprocessing import StandardScaler | |
| BACKEND_ROOT = Path(__file__).resolve().parents[1] | |
| if str(BACKEND_ROOT) not in sys.path: | |
| sys.path.insert(0, str(BACKEND_ROOT)) | |
| SCRIPT_ROOT = Path(__file__).resolve().parent | |
| if str(SCRIPT_ROOT) not in sys.path: | |
| sys.path.insert(0, str(SCRIPT_ROOT)) | |
| from app.config import DEFAULT_RUNTIME_REFINEMENT_REPORT_PATH, DEFAULT_RUNTIME_REFINER_PATH | |
| from app.ml.runtime_refinement import RuntimeScreeningRefiner | |
| from app.services.image_quality import ImageQualityService | |
| from app.services.prediction import ScreeningPredictor | |
| from train_efficientnet import ( | |
| ARCHIVE_ROOT, | |
| _balanced_group_split, | |
| _build_records, | |
| _load_image_with_fallback, | |
| ) | |
| def _metric_block(labels: np.ndarray, predictions: np.ndarray) -> dict[str, float]: | |
| return { | |
| "accuracy": round(float(accuracy_score(labels, predictions)), 4), | |
| "precision": round(float(precision_score(labels, predictions, zero_division=0)), 4), | |
| "recall": round(float(recall_score(labels, predictions, zero_division=0)), 4), | |
| "f1": round(float(f1_score(labels, predictions, zero_division=0)), 4), | |
| } | |
| def _build_dataset(records) -> tuple[np.ndarray, np.ndarray, np.ndarray]: | |
| quality_service = ImageQualityService() | |
| predictor = ScreeningPredictor() | |
| predictor.runtime_screening_refiner = None | |
| predictor._runtime_screening_refiner_load_attempted = True | |
| feature_rows: list[list[float]] = [] | |
| labels: list[int] = [] | |
| base_predictions: list[int] = [] | |
| for record in records: | |
| with record.image_path.open("rb") as handle: | |
| quality, processed = quality_service.evaluate(handle.read()) | |
| prediction = predictor.predict(processed, quality) if quality.passed else None | |
| if prediction is None and quality_service.allows_raw_frame_rescue(quality): | |
| raw_image = _load_image_with_fallback(record.image_path).convert("RGB") | |
| raw_prediction = predictor.predict(raw_image, quality) | |
| if predictor.should_accept_raw_frame_rescue(raw_prediction): | |
| quality = quality_service.build_raw_frame_rescue_assessment(quality) | |
| prediction = raw_prediction | |
| if prediction is None: | |
| base_risk = 0.0 | |
| uncertainty = 1.0 | |
| predicted_hemoglobin = None | |
| base_likely = False | |
| base_prediction = 0 | |
| else: | |
| base_risk = float( | |
| prediction.confidence_breakdown.get("raw_anemia_risk", prediction.anemia_risk) | |
| ) | |
| uncertainty = float(prediction.uncertainty) | |
| predicted_hemoglobin = prediction.predicted_hemoglobin | |
| base_likely = str( | |
| prediction.confidence_breakdown.get("base_screening_label", prediction.screening_label) | |
| ) == "anemia_likely" | |
| base_prediction = int(prediction.screening_label == "anemia_likely") | |
| feature_rows.append( | |
| RuntimeScreeningRefiner()._feature_vector( | |
| base_anemia_risk=base_risk, | |
| uncertainty=uncertainty, | |
| predicted_hemoglobin=predicted_hemoglobin, | |
| quality=quality, | |
| base_likely=base_likely, | |
| ) | |
| ) | |
| labels.append(int(record.label)) | |
| base_predictions.append(base_prediction) | |
| return ( | |
| np.asarray(feature_rows, dtype=np.float32), | |
| np.asarray(labels, dtype=np.int32), | |
| np.asarray(base_predictions, dtype=np.int32), | |
| ) | |
| def _evaluate_deployed_records(records, *, use_refiner: bool) -> dict[str, float]: | |
| quality_service = ImageQualityService() | |
| predictor = ScreeningPredictor() | |
| if not use_refiner: | |
| predictor.runtime_screening_refiner = None | |
| predictor._runtime_screening_refiner_load_attempted = True | |
| labels: list[int] = [] | |
| predictions: list[int] = [] | |
| for record in records: | |
| with record.image_path.open("rb") as handle: | |
| quality, processed = quality_service.evaluate(handle.read()) | |
| prediction = predictor.predict(processed, quality) if quality.passed else None | |
| if prediction is None and quality_service.allows_raw_frame_rescue(quality): | |
| raw_image = _load_image_with_fallback(record.image_path).convert("RGB") | |
| raw_prediction = predictor.predict(raw_image, quality) | |
| if predictor.should_accept_raw_frame_rescue(raw_prediction): | |
| quality = quality_service.build_raw_frame_rescue_assessment(quality) | |
| prediction = raw_prediction | |
| labels.append(int(record.label)) | |
| predictions.append(int(prediction is not None and prediction.screening_label == "anemia_likely")) | |
| return _metric_block( | |
| np.asarray(labels, dtype=np.int32), | |
| np.asarray(predictions, dtype=np.int32), | |
| ) | |
| def _choose_threshold(labels: np.ndarray, probabilities: np.ndarray) -> tuple[float, dict[str, float]]: | |
| best_threshold = 0.5 | |
| best_metrics: dict[str, float] | None = None | |
| for threshold in np.linspace(0.3, 0.7, 41): | |
| predictions = (probabilities >= threshold).astype(np.int32) | |
| metrics = _metric_block(labels, predictions) | |
| if best_metrics is None or metrics["f1"] > best_metrics["f1"] or ( | |
| metrics["f1"] == best_metrics["f1"] and metrics["precision"] > best_metrics["precision"] | |
| ): | |
| best_threshold = float(threshold) | |
| best_metrics = metrics | |
| assert best_metrics is not None | |
| return best_threshold, best_metrics | |
| def main() -> None: | |
| records = _build_records(ARCHIVE_ROOT) | |
| if not records: | |
| raise RuntimeError(f"No evaluation records found in {ARCHIVE_ROOT}.") | |
| train_records, val_records = _balanced_group_split(records, test_size=0.2, n_splits=32) | |
| train_roi = [record for record in train_records if record.source == "roi_original"] | |
| val_roi = [record for record in val_records if record.source == "roi_original"] | |
| X_train, y_train, _ = _build_dataset(train_roi) | |
| X_val, y_val, base_predictions = _build_dataset(val_roi) | |
| model = Pipeline( | |
| [ | |
| ("scaler", StandardScaler()), | |
| ( | |
| "logreg", | |
| LogisticRegression( | |
| C=0.3, | |
| max_iter=4000, | |
| class_weight="balanced", | |
| random_state=42, | |
| ), | |
| ), | |
| ] | |
| ) | |
| model.fit(X_train, y_train) | |
| probabilities = model.predict_proba(X_val)[:, 1] | |
| selected_threshold, stage_metrics_after = _choose_threshold(y_val, probabilities) | |
| metrics_before = _evaluate_deployed_records(val_roi, use_refiner=False) | |
| refiner = RuntimeScreeningRefiner( | |
| model=model, | |
| threshold=round(selected_threshold, 4), | |
| report={ | |
| "validation_size": int(len(y_val)), | |
| "metrics_before": metrics_before, | |
| "selected_threshold": round(selected_threshold, 4), | |
| }, | |
| ) | |
| refiner.save(DEFAULT_RUNTIME_REFINER_PATH) | |
| metrics_after = _evaluate_deployed_records(val_roi, use_refiner=True) | |
| report = { | |
| "version": refiner.version, | |
| "method": refiner.method, | |
| "validation_size": int(len(y_val)), | |
| "selected_threshold": round(selected_threshold, 4), | |
| "metrics_before": metrics_before, | |
| "metrics_after": metrics_after, | |
| "stage_metrics_after": stage_metrics_after, | |
| } | |
| DEFAULT_RUNTIME_REFINEMENT_REPORT_PATH.write_text(json.dumps(report, indent=2), encoding="utf-8") | |
| print("\nRuntime screening refinement metrics") | |
| print(f"validation_size: {report['validation_size']}") | |
| print(f"selected_threshold: {report['selected_threshold']:.4f}") | |
| print("before:", report["metrics_before"]) | |
| print("after:", report["metrics_after"]) | |
| if __name__ == "__main__": | |
| main() | |