tta segmentation #ml #pytorch #python

tta_segmentation.py

#https://github.com/qubvel/ttach
import ttach as tta

transforms = tta.Compose(
  [
    tta.HorizontalFlip(),
    tta.VerticalFlip(),
    #                 tta.Rotate90(angles=[0, 180]),
    #                 tta.Scale(scales=[1, 2, 4]),
    #                 tta.Multiply(factors=[0.9, 1, 1.1]),        
  ]
)

preds = []
for transformer in transforms:
  #print(1)
  augmented_images = transformer.augment_image(images)
  model_output = model(augmented_images)
  deaug_mask = transformer.deaugment_mask(model_output)
  preds.append(deaug_mask)
preds = torch.mean(torch.stack(preds), dim=0)