main.py

#!/usr/bin/env python3
import argparse
import sys
import subprocess
import os
import json
import pandas as pd

from src.tag2pixel.utils import load_config, setup_logging
from src.tag2pixel.osm import run_osm_extraction
from src.tag2pixel.labeling import run_label_generation
from src.tag2pixel.features import run_feature_extraction

logger = setup_logging("Orchestrator")

def run_r_training(cfg):
    logger.info("--- Running: Model Training (R) ---")
    try:
        # Pass the config path if needed, though R script currently assumes root or fixed paths
        subprocess.run(["Rscript", "src/r_scripts/train_rf.R"], check=True)
    except subprocess.CalledProcessError as e:
        logger.error(f"❌ Error during R training: {e}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description="Tag2Pixel: Universal Geospatial Framework")
    parser.add_argument("--config", default="config/config.yaml", help="Path to config file")
    parser.add_argument("--step", choices=["osm", "labels", "features", "train", "all"], default="all")
    args = parser.parse_args()

    try:
        cfg = load_config(args.config)
    except Exception as e:
        logger.error(f"❌ Failed to load config: {e}")
        sys.exit(1)

    # Ensure directories
    for d in ["data/raw_osm_data", "data/ground_truth", "data/labels", "data/training", "output/plots"]:
        os.makedirs(d, exist_ok=True)

    steps = [args.step] if args.step != "all" else ["osm", "labels", "features", "train"]

    if "osm" in steps:
        logger.info("--- Step 1/4: OSM Extraction ---")
        run_osm_extraction(cfg)

    if "labels" in steps:
        logger.info("--- Step 2/4: Label Generation ---")
        run_label_generation(cfg)

    if "features" in steps:
        logger.info("--- Step 3/4: Feature Extraction ---")
        run_feature_extraction(cfg)

    if "train" in steps:
        logger.info("--- Step 4/4: Model Training ---")
        run_r_training(cfg)
        
        # Post-training validation
        if os.path.exists("output/metrics.json"):
            with open("output/metrics.json", "r") as f:
                m = json.load(f)
            logger.info(f"✨ Results - AUC: {m.get('auc',0):.4f}, Acc: {m.get('accuracy',0):.4f}")

    logger.info("🎉 Tag2Pixel Pipeline Finished!")

if __name__ == "__main__":
    main()