How to generate geotif labels from a image with a trained model?

[lcoandrade]

Hi there!

I'm working on a Kaggle notebook with a custom dataset based on NAIP images and Chesapeake labels.

I want to perform a segmentation task.

Right now, I'm able to train with one GPU. How can I, from this trained model, generate geotif labels from a image?

[adamjstewart]

@adriantre is working on a way to automatically do this in SemanticSegmentationTask in #1407. If you just want a PNG, it's really easy to do. If you want a GeoTIFF, we have a little bit of work required to store the CRS/transform needed to save geospatial metadata. You can also do this outside of Lightning by loading the model and data loader and iterating over inputs, making predictions, and saving to disk using the same rasterio profile.

[calebrob6]

In cell [16] I would try dividing image by 255.0

In GeoDataModule augmentation is done when batches are transfered to the GPU, but in cell [16] you are bypassing this by creating a dataloader from the datamodule's test dataset so I don't think augmentation is being applied.

Also, something is still wrong with the images... maybe just try image[:,:,:3] instead of image[:,:,[2,1,0]]?

[calebrob6]

More generally, if your model output is somewhat constant, or definitely not what you expect, a good first debugging step is to examine the inputs closely and see if they are aligning with the type of inputs that the model was trained on. E.g. if the model was trained on inputs where you have done standardization per-channel, it is important to do this same step at inference time as well.

[lcoandrade]

I understand, but I'm treating the images used to train and those used in the prediction the same way. I just create a dataset with the NAIP images and Chesapeake labels and make a random split to generate train, validation and test datasets.

Custom datasets:

class NAIPImages(RasterDataset):
    filename_glob = "m_*.tif"
    is_image = True
    separate_files = False
    
class ChesapeakeLabels(RasterDataset):
    filename_glob = "m_*.tif"
    is_image = False
    separate_files = False

Making them individually:

naip_root = os.path.join(INPUT_DIR, 'naip_images')
naip_images = NAIPImages(naip_root)
print(naip_images)

chesapeake_root = os.path.join(INPUT_DIR, "chesapeake_labels")
chesapeake_labels = ChesapeakeLabels(chesapeake_root)
print(chesapeake_labels)

Creating an intersection:

dataset = naip_images & chesapeake_labels # this means I'm creating an IntersectionDataset

From this, I create my custom datamodule that is used to train my network:

class CustomGeoDataModule(GeoDataModule):
    def setup(self, stage: str) -> None:
        """Set up datasets.

        Args:
            stage: Either 'fit', 'validate', 'test', or 'predict'.
        """
        self.dataset = self.dataset_class(**self.kwargs)
        
        generator = torch.Generator().manual_seed(0)
        (
            self.train_dataset,
            self.val_dataset,
            self.test_dataset,
        ) = random_bbox_assignment(dataset, [0.6, 0.2, 0.2], generator)
        
        if stage in ["fit"]:
            self.train_batch_sampler = RandomBatchGeoSampler(
                self.train_dataset, self.patch_size, self.batch_size, self.length
            )
        if stage in ["fit", "validate"]:
            self.val_sampler = GridGeoSampler(
                self.val_dataset, self.patch_size, self.patch_size
            )
        if stage in ["test"]:
            self.test_sampler = GridGeoSampler(
                self.test_dataset, self.patch_size, self.patch_size
            )
            
datamodule = CustomGeoDataModule(
    dataset_class = type(dataset), # GeoDataModule kwargs
    batch_size = 10, # GeoDataModule kwargs
    patch_size = IMG_SIZE, # GeoDataModule kwargs
    length = SAMPLE_SIZE, # GeoDataModule kwargs
    num_workers = WORKERS, # GeoDataModule kwargs
    dataset1 = naip_images, # IntersectionDataset kwargs
    dataset2 = chesapeake_labels, # IntersectionDataset kwargs
    collate_fn = stack_samples, # IntersectionDataset kwargs
)

This datamodule is used to train and to predict.

[calebrob6]

I know, but when you do:

trainer.test(
    model=task,
    datamodule=datamodule
)

I'm pretttty sure the images are being divided by 255

and when you do:

test_dataset = datamodule.test_dataset
test_sampler = RandomGeoSampler(test_dataset, size=IMG_SIZE, length=SAMPLE_SIZE)
test_dataloader = DataLoader(test_dataset, sampler=test_sampler, collate_fn=stack_samples)

n = 0
for sample in test_dataloader:
    ...

I'm pretty sure the images are not being divided by 255.

[lcoandrade]

Ahhhh, I'll insert an:

image = image/255.

To see what will happen.

[calebrob6]

Also change up how the images are being visualized ;)

[lcoandrade]

Just making:

image = image/255.

Solved the problem!

Thank you very much!

Now I would like to proceed to generate a geotiff for a whole image.

[lcoandrade]

Version 10 is being commit to show the correct outputs!