1import click
2import lightning as L
3from methane_super_emitters.model import SuperEmitterDetector
4from methane_super_emitters.datamodule import TROPOMISuperEmitterDataModule
5
6
7@click.command()
8@click.option("-i", "--input-dir", help="Data directory")
9@click.option("-m", "--max-epochs", help="Maximum number of epochs", default=1)
10def train_model(input_dir, max_epochs):
11 fields = ["methane", "u10", "v10", "qa"]
12 model = SuperEmitterDetector(fields=fields)
13 datamodule = TROPOMISuperEmitterDataModule(input_dir, fields=fields)
14 trainer = L.Trainer(max_epochs=max_epochs)
15 trainer.fit(model=model, datamodule=datamodule)
16 trainer.test(model=model, datamodule=datamodule)
17
18
19if __name__ == "__main__":
20 train_model()