from deep_learning.data.oxford_pets import SegmentationDataSource
from deep_learning.env.resolve import resolve_env, resolve_path, resolve_saved
from deep_learning.models.segmentation import SegmentationModelBuilder
from deep_learning.pipeline import (
    SupervisedModelPipeline,
    PipelineRunner
)
from deep_learning.pipeline.specs.configs import CheckpointConfig, CheckpointLoadRules, TrainingRule


pipeline = resolve_env(
    # 开发配置
    SupervisedModelPipeline(
        name="segmentation",
        data_source=SegmentationDataSource(
            images_path=resolve_path("data/dev/oxford_pets/images"),
            annotations_path=resolve_path("data/dev/oxford_pets/annotations/trimaps"),
            image_size=(200, 200),
            batch_size=2,
            validation_batches=1,
            example_count=5,
            example_output_dir=resolve_path("local/examples/segmentation")
        ),
        model_builder=SegmentationModelBuilder(
            image_size=(200, 200),
            num_classes=3,
            model_filters=(8,)
        ),
        training_rule=TrainingRule(
            epochs=1,
            steps_per_epoch=None
        )
    ),

    # 生产配置
    SupervisedModelPipeline(
        name="segmentation",
        data_source=SegmentationDataSource(
            images_path=resolve_path("~/.keras/datasets/vgg_perts_images_extracted/images"),
            annotations_path=resolve_path("~/.keras/datasets/vgg_pets_annotations_extracted/annotations/trimaps"),
            image_size=(200, 200),
            batch_size=64,
            validation_batches=15,
            example_count=5,
            example_output_dir=resolve_path("local/examples/segmentation")
        ),
        model_builder=SegmentationModelBuilder(
            image_size=(200, 200),
            num_classes=3,
            model_filters=(64, 128, 256)
        ),
        training_rule=TrainingRule(
            epochs=50,
            steps_per_epoch=None
        ),
        checkpoint_load_rules=CheckpointLoadRules(
            export=CheckpointConfig(epoch=26),
            test=CheckpointConfig(dirs=[resolve_saved("models/segmentation")], suffix=".keras")
        )
    )
)

pipeline_runner = PipelineRunner(pipeline)


if __name__ == "__main__":
    pipeline_runner()