def test_predict_numpy():
img = np.ones((1, 3, 10, 20))
model = SemanticSegmentation(2)
data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess())
out = model.predict(img, data_source="numpy", data_pipeline=data_pipe)
assert isinstance(out[0], torch.Tensor)
assert out[0].shape == (196, 196)
labels_map = SegmentationLabels.create_random_labels_map(num_classes=21)
datamodule.set_labels_map(labels_map)
datamodule.show_train_batch(["load_sample", "post_tensor_transform"])
# 3. Build the model
model = SemanticSegmentation(
backbone="torchvision/fcn_resnet50",
num_classes=21,
)
# 4. Create the trainer.
trainer = flash.Trainer(
max_epochs=1,
fast_dev_run=1,
)
# 5. Train the model
trainer.finetune(model, datamodule=datamodule, strategy="freeze")
# 6. Predict what's on a few images!
model.serializer = SegmentationLabels(labels_map, visualize=True)
predictions = model.predict([
"data/CameraRGB/F61-1.png",
"data/CameraRGB/F62-1.png",
"data/CameraRGB/F63-1.png",
])
# 7. Save it!
trainer.save_checkpoint("semantic_segmentation_model.pt")