Staging-model-eval

applications/DynaCLR/evaluation/linear_classifiers/apply_linear_classifier.py

@@ -61,6 +61,19 @@ def format_predictions_markdown(adata, task: str) -> str:
         lines.append(f"**Classes:** {', '.join(adata.uns[classes_key])}")
         lines.append("")
 
+    artifact_key = f"classifier_{task}_artifact"
+    if artifact_key in adata.uns.keys():
+        lines.append("### Classifier Provenance")
+        lines.append("")
+        lines.append(f"- **Artifact:** {adata.uns[artifact_key]}")
+        id_key = f"classifier_{task}_id"
+        if id_key in adata.uns.keys():
+            lines.append(f"- **Artifact ID:** {adata.uns[id_key]}")
+        version_key = f"classifier_{task}_version"
+        if version_key in adata.uns.keys():
+            lines.append(f"- **Artifact Version:** {adata.uns[version_key]}")
+        lines.append("")
+
     return "\n".join(lines)
 
 
@@ -88,36 +101,52 @@ def main(config: Path):
         click.echo(f"\n❌ Failed to load configuration: {e}", err=True)
         raise click.Abort()
 
+    write_path = (
+        Path(inference_config.output_path)
+        if inference_config.output_path is not None
+        else Path(inference_config.embeddings_path)
+    )
+
     click.echo(f"\n✓ Configuration loaded: {config}")
     click.echo(f"  Model: {inference_config.model_name}")
     click.echo(f"  Version: {inference_config.version}")
     click.echo(f"  Embeddings: {inference_config.embeddings_path}")
-    click.echo(f"  Output: {inference_config.output_path}")
+    click.echo(f"  Output: {write_path}")
 
     try:
-        pipeline, loaded_config = load_pipeline_from_wandb(
+        pipeline, loaded_config, artifact_metadata = load_pipeline_from_wandb(
             wandb_project=inference_config.wandb_project,
             model_name=inference_config.model_name,
             version=inference_config.version,
             wandb_entity=inference_config.wandb_entity,
         )
 
         task = loaded_config["task"]
+        marker = loaded_config.get("marker")
+        task_key = f"{task}_{marker}" if marker else task
 
         click.echo(f"\nLoading embeddings from: {inference_config.embeddings_path}")
         adata = read_zarr(inference_config.embeddings_path)
         click.echo(f"✓ Loaded embeddings: {adata.shape}")
 
-        adata = predict_with_classifier(adata, pipeline, task)
+        if inference_config.include_wells:
+            click.echo(f"  Well filter: {inference_config.include_wells}")
+
+        adata = predict_with_classifier(
+            adata,
+            pipeline,
+            task_key,
+            artifact_metadata=artifact_metadata,
+            include_wells=inference_config.include_wells,
+        )
 
-        output_path = Path(inference_config.output_path)
-        output_path.parent.mkdir(parents=True, exist_ok=True)
+        write_path.parent.mkdir(parents=True, exist_ok=True)
 
-        click.echo(f"\nSaving predictions to: {output_path}")
-        adata.write_zarr(output_path)
+        click.echo(f"\nSaving predictions to: {write_path}")
+        adata.write_zarr(write_path)
         click.echo("✓ Saved predictions")
 
-        click.echo("\n" + format_predictions_markdown(adata, task))
+        click.echo("\n" + format_predictions_markdown(adata, task_key))
 
         click.echo("\n✓ Inference complete!")
 

applications/DynaCLR/evaluation/linear_classifiers/configs/example_linear_classifier_inference.yaml

@@ -21,8 +21,16 @@ wandb_entity: null
 # Path to embeddings zarr file for inference
 embeddings_path: /path/to/embeddings.zarr
 
-# Path to save output zarr file with predictions
-output_path: /path/to/output_with_predictions.zarr
+# Path to save output zarr file with predictions.
+# When omitted (or null), predictions are written back to embeddings_path.
+# output_path: /path/to/output_with_predictions.zarr
 
-# Whether to overwrite output if it already exists
+# Well prefixes to restrict predictions to (optional).
+# When omitted, all cells are predicted. Cells in other wells get NaN.
+# Useful for organelle-specific classifiers where different wells have different markers.
+# include_wells:
+#   - A/1
+#   - A/2
+
+# Whether to overwrite output if it already exists (only used when output_path is set)
 overwrite: false