def test_predict_numpy(): img = np.ones((1, 3, 64, 64)) model = SemanticSegmentation(2, backbone="mobilenetv3_large_100") data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess(num_classes=1)) out = model.predict(img, data_source="numpy", data_pipeline=data_pipe) assert isinstance(out[0], list) assert len(out[0]) == 64 assert len(out[0][0]) == 64
def test_predict_numpy(): img = np.ones((1, 3, 10, 20)) model = SemanticSegmentation(2) data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess( num_classes=1)) out = model.predict(img, data_source="numpy", data_pipeline=data_pipe) assert isinstance(out[0], torch.Tensor) assert out[0].shape == (10, 20)
def test_predict_numpy(): img = np.ones((1, 3, 10, 20)) model = SemanticSegmentation(2) data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess( num_classes=1)) out = model.predict(img, data_source="numpy", data_pipeline=data_pipe) assert isinstance(out[0], list) assert len(out[0]) == 10 assert len(out[0][0]) == 20
def test_predict_tensor(): img = torch.rand(1, 3, 10, 20) model = SemanticSegmentation(2) data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess( num_classes=1)) out = model.predict(img, data_source="tensors", data_pipeline=data_pipe) assert isinstance(out[0], list) assert len(out[0]) == 10 assert len(out[0][0]) == 20
def test_jit(tmpdir, jitter, args): path = os.path.join(tmpdir, "test.pt") model = SemanticSegmentation(2) model.eval() model = jitter(model, *args) torch.jit.save(model, path) model = torch.jit.load(path) out = model(torch.rand(1, 3, 32, 32)) assert isinstance(out, torch.Tensor) assert out.shape == torch.Size([1, 2, 32, 32])
def test_predict_numpy(): img = np.ones((1, 3, 64, 64)) model = SemanticSegmentation(2, backbone="mobilenetv3_large_100") datamodule = SemanticSegmentationData.from_numpy(predict_data=img, batch_size=1) trainer = Trainer() out = trainer.predict(model, datamodule=datamodule, output="labels") assert isinstance(out[0][0], list) assert len(out[0][0]) == 64 assert len(out[0][0][0]) == 64
def test_forward(num_classes, img_shape): model = SemanticSegmentation( num_classes=num_classes, backbone='torchvision/fcn_resnet50', ) B, C, H, W = img_shape img = torch.rand(B, C, H, W) out = model(img) assert out.shape == (B, num_classes, H, W)
def test_forward(num_classes, img_shape): model = SemanticSegmentation( num_classes=num_classes, backbone="resnet50", head="fpn", ) B, C, H, W = img_shape img = torch.rand(B, C, H, W) out = model(img) assert out.shape == (B, num_classes, H, W)
val_split=0.1, image_size=(200, 200), num_classes=21, ) # 2.2 Visualise the samples datamodule.show_train_batch(["load_sample", "post_tensor_transform"]) # 3.a List available backbones and heads print(f"Backbones: {SemanticSegmentation.available_backbones()}") print(f"Heads: {SemanticSegmentation.available_heads()}") # 3.b Build the model model = SemanticSegmentation( backbone="mobilenet_v3_large", head="fcn", num_classes=datamodule.num_classes, serializer=SegmentationLabels(visualize=False), ) # 4. Create the trainer. trainer = flash.Trainer(fast_dev_run=True) # 5. Train the model trainer.finetune(model, datamodule=datamodule, strategy="freeze") # 6. Segment a few images! predictions = model.predict([ "data/CameraRGB/F61-1.png", "data/CameraRGB/F62-1.png", "data/CameraRGB/F63-1.png", ])
# 2.1 Load the data datamodule = SemanticSegmentationData.from_folders( train_folder="data/CameraRGB", train_target_folder="data/CameraSeg", batch_size=4, val_split=0.3, image_size=(200, 200), # (600, 800) num_classes=21, ) # 2.2 Visualise the samples datamodule.show_train_batch(["load_sample", "post_tensor_transform"]) # 3. Build the model model = SemanticSegmentation(backbone="torchvision/fcn_resnet50", num_classes=datamodule.num_classes, serializer=SegmentationLabels(visualize=True)) # 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") predictions = model.predict([ "data/CameraRGB/F61-1.png", "data/CameraRGB/F62-1.png", "data/CameraRGB/F63-1.png",
def test_freeze(): model = SemanticSegmentation(2) model.freeze() for p in model.backbone.parameters(): assert p.requires_grad is False
def test_non_existent_backbone(): with pytest.raises(KeyError): SemanticSegmentation(2, "i am never going to implement this lol")
def test_init_train(tmpdir): model = SemanticSegmentation(num_classes=10) train_dl = torch.utils.data.DataLoader(DummyDataset()) trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True) trainer.finetune(model, train_dl, strategy="freeze_unfreeze")
def test_smoke(): model = SemanticSegmentation(num_classes=1) assert model is not None
def test_available_pretrained_weights(): assert SemanticSegmentation.available_pretrained_weights("resnet18") == [ "imagenet", "ssl", "swsl" ]
def test_serve(): model = SemanticSegmentation(2) # TODO: Currently only servable once a preprocess has been attached model._preprocess = SemanticSegmentationPreprocess() model.eval() model.serve()
"https://github.com/ongchinkiat/LyftPerceptionChallenge/releases/download/v0.1/carla-capture-20180513A.zip", "./data", ) datamodule = SemanticSegmentationData.from_folders( train_folder="data/CameraRGB", train_target_folder="data/CameraSeg", val_split=0.1, image_size=(256, 256), num_classes=21, ) # 2. Build the task model = SemanticSegmentation( backbone="mobilenetv3_large_100", head="fpn", num_classes=datamodule.num_classes, ) # 3. Create the trainer and finetune the model trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count()) trainer.finetune(model, datamodule=datamodule, strategy="freeze") # 4. Segment a few images! predictions = model.predict([ "data/CameraRGB/F61-1.png", "data/CameraRGB/F62-1.png", "data/CameraRGB/F63-1.png", ]) print(predictions)
"./data", ) datamodule = SemanticSegmentationData.from_folders( train_folder="data/CameraRGB", train_target_folder="data/CameraSeg", val_split=0.1, transform_kwargs=dict(image_size=(256, 256)), num_classes=21, batch_size=4, ) # 2. Build the task model = SemanticSegmentation( backbone="mobilenetv3_large_100", head="fpn", num_classes=datamodule.num_classes, ) # 3. Create the trainer and finetune the model trainer = flash.Trainer(max_epochs=3, gpus=torch.cuda.device_count()) trainer.finetune(model, datamodule=datamodule, strategy="freeze") # 4. Segment a few images! datamodule = SemanticSegmentationData.from_files( predict_files=[ "data/CameraRGB/F61-1.png", "data/CameraRGB/F62-1.png", "data/CameraRGB/F63-1.png", ], batch_size=3,
# Copyright The PyTorch Lightning team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from flash.image import SemanticSegmentation model = SemanticSegmentation.load_from_checkpoint( "https://flash-weights.s3.amazonaws.com/semantic_segmentation_model.pt") model.serve()
def test_load_from_checkpoint_dependency_error(): with pytest.raises(ModuleNotFoundError, match=re.escape("'lightning-flash[image]'")): SemanticSegmentation.load_from_checkpoint("not_a_real_checkpoint.pt")
def test_map_labels(tmpdir): tmp_dir = Path(tmpdir) # create random dummy data images = [ str(tmp_dir / "img1.png"), str(tmp_dir / "img2.png"), str(tmp_dir / "img3.png"), ] targets = [ str(tmp_dir / "labels_img1.png"), str(tmp_dir / "labels_img2.png"), str(tmp_dir / "labels_img3.png"), ] labels_map: Dict[int, Tuple[int, int, int]] = { 0: [0, 0, 0], 1: [255, 255, 255], } num_classes: int = len(labels_map.keys()) img_size: Tuple[int, int] = (128, 128) create_random_data(images, targets, img_size, num_classes) # instantiate the data module dm = SemanticSegmentationData.from_files( train_files=images, train_targets=targets, val_files=images, val_targets=targets, batch_size=2, num_workers=0, num_classes=num_classes, ) assert dm is not None assert dm.train_dataloader() is not None # disable visualisation for testing assert dm.data_fetcher.block_viz_window is True dm.set_block_viz_window(False) assert dm.data_fetcher.block_viz_window is False dm.show_train_batch("load_sample") dm.show_train_batch("to_tensor_transform") # check training data data = next(iter(dm.train_dataloader())) imgs, labels = data[DefaultDataKeys.INPUT], data[DefaultDataKeys.TARGET] assert imgs.shape == (2, 3, 128, 128) assert labels.shape == (2, 128, 128) assert labels.min().item() == 0 assert labels.max().item() == 1 assert labels.dtype == torch.int64 # now train with `fast_dev_run` model = SemanticSegmentation(num_classes=2, backbone="resnet50", head="fpn") trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True) trainer.finetune(model, dm, strategy="freeze_unfreeze")
def test_serve(): model = SemanticSegmentation(2) model.eval() model.serve()