How to augment the CropModel during training.

[bw4sz]

@martibosch mentioned wanting custom transformers for the CropModel.

One approach would be to subclass the model and override the transform.

from torchvision import transforms
from deepforest.model import CropModel
class MyCropModel(CropModel):
    def get_transform(self, augmentations=None):
        resize_dims = self.config["cropmodel"].get("resize", [224, 224])
        t = [transforms.ToTensor()]
        norm = self.normalize()
        if norm is not None:
            t.append(norm)
        t.append(transforms.Resize(resize_dims))
        if augmentations is not None:
            t.extend([
                transforms.RandomHorizontalFlip(0.5),
                transforms.RandomRotation(10),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            ])
        return transforms.Compose(t)
crop_model = MyCropModel(num_classes=2)
crop_model.load_from_disk(train_dir="...", val_dir="...")
crop_model.create_trainer(...)
crop_model.trainer.fit(crop_model)

Is this sufficient, if so should we add to docs? Or should be make this an arg for the CropModel(...,transform=)?