from deep_learning.data.coco import YoloDataSource
from deep_learning.env.resolve import resolve_env, resolve_path, resolve_saved
from deep_learning.models.yolo import YoloModelBuilder
from deep_learning.pipeline import (
    SupervisedModelPipeline,
    PipelineRunner
)
from deep_learning.pipeline.specs.configs import CheckpointConfig, CheckpointLoadRules, TrainingRule


pipeline = resolve_env(
    # 开发配置
    SupervisedModelPipeline(
        name="yolo",
        data_source=YoloDataSource(
            images_path=resolve_path("data/dev/coco/train2017"),
            annotation_file=resolve_path("data/dev/coco/annotations/instances_train2017.json"),
            image_size=448,
            grid_size=6,
            batch_size=2,
            validation_batches=1,
            max_objects_per_image=4,
            example_count=5,
            example_output_dir=resolve_path("local/examples/yolo")
        ),
        model_builder=YoloModelBuilder(
            image_size=448,
            grid_size=6,
            num_labels=91,
            backbone_preset="resnet_50_imagenet"
        ),
        training_rule=TrainingRule(
            epochs=1,
            steps_per_epoch=None
        )
    ),

    # 生产配置
    SupervisedModelPipeline(
        name="yolo",
        data_source=YoloDataSource(
            images_path=resolve_path("~/data/coco/train2017"),
            annotation_file=resolve_path("~/data/coco/annotations/instances_train2017.json"),
            image_size=448,
            grid_size=6,
            batch_size=32,
            validation_batches=500,
            max_objects_per_image=4,
            example_count=5,
            example_output_dir=resolve_path("local/examples/yolo")
        ),
        model_builder=YoloModelBuilder(
            image_size=448,
            grid_size=6,
            num_labels=91,
            backbone_preset="resnet_50_imagenet"
        ),
        training_rule=TrainingRule(
            epochs=100,
            steps_per_epoch=None
        ),
        checkpoint_load_rules=CheckpointLoadRules(
            export=CheckpointConfig(epoch=18),
            test=CheckpointConfig(dirs=[resolve_saved("models/yolo")], suffix=".keras")
        )
    )
)

pipeline_runner = PipelineRunner(pipeline)


if __name__ == "__main__":
    pipeline_runner()