feat(validation): add visualization output module

This commit is contained in:
2026-03-19 10:29:13 +08:00
parent 1e36e52112
commit 4e869bb693

View File

@@ -0,0 +1,74 @@
from __future__ import annotations
from pathlib import Path
from rdkit import Chem
from macro_lactone_toolkit.visualization import save_fragment_png, save_numbered_molecule_png
def get_output_paths(output_dir: Path, source_id: str, ring_size: int, classification: str) -> dict:
"""Get organized output paths for a molecule."""
ring_dir = output_dir / f"ring_size_{ring_size}"
if classification == "standard_macrolactone":
base_dir = ring_dir / "standard"
numbered_dir = base_dir / "numbered"
sidechains_dir = base_dir / "sidechains" / source_id
elif classification == "non_standard_macrocycle":
base_dir = ring_dir / "non_standard" / "original"
numbered_dir = base_dir
sidechains_dir = None
else:
base_dir = ring_dir / "rejected" / "original"
numbered_dir = base_dir
sidechains_dir = None
numbered_dir.mkdir(parents=True, exist_ok=True)
if sidechains_dir:
sidechains_dir.mkdir(parents=True, exist_ok=True)
return {
"numbered_image": numbered_dir / f"{source_id}_numbered.png",
"sidechains_dir": sidechains_dir,
}
def save_numbered_molecule(
smiles: str,
output_path: Path,
ring_size: int | None = None,
size: tuple[int, int] = (800, 800),
) -> Path | None:
"""Save numbered molecule image."""
try:
return save_numbered_molecule_png(
smiles,
output_path,
ring_size=ring_size,
size=size,
)
except Exception as e:
print(f"Failed to save numbered image: {e}")
return None
def save_fragment_images(
fragments: list,
output_dir: Path,
source_id: str,
size: tuple[int, int] = (400, 400),
) -> list[str]:
"""Save fragment images and return paths."""
paths = []
for i, fragment in enumerate(fragments):
try:
output_path = output_dir / f"{source_id}_frag_{i}_pos{fragment.cleavage_position}.png"
save_fragment_png(fragment.fragment_smiles_plain, output_path, size=size)
paths.append(str(output_path.relative_to(output_dir.parent.parent)))
except Exception as e:
print(f"Failed to save fragment {i}: {e}")
paths.append(None)
return paths