segsem.py

# Load the data
datamodule = SemanticSegmentationData.from_folders(
    train_folder="./train",
    train_target_folder="./targets",
    num_classes=21)

# Build the model
model = SemanticSegmentation(
    backbone="torchvision/fcn_resnet50",
    num_classes=datamodule.num_classes,
    serializer=SegmentationLabels(visualize=True))

# Create the trainer
trainer = flash.Trainer()

# Train the model
trainer.finetune(model, datamodule=datamodule)

# Make a prediction
model.predict(["image.png"])