Coverage for cafl4ds/ssl/augment.py: 100%

1"""Augmentation pipelines for the SSL methods. 

2 

3Joint-embedding methods (SimSiam here, BYOL/SimCLR later) learn by pulling together two 

4*augmented views* of the same image, so :class:`TwoView` returns an independent pair. MAE 

5needs only light spatial augmentation (its learning signal comes from masking), so it uses a 

6single view. 

7 

8Simplification for the Phase-0 harness: a transform draws its random parameters **once per 

9call** and applies them to the whole batch (fast on CPU). Two views therefore differ from 

10each other (two independent draws) but samples within a view share a crop/jitter. That is 

11enough to exercise the joint-embedding mechanics; per-sample augmentation can replace this 

12without touching the method code. 

13""" 

14 

15from __future__ import annotations 

16 

17import torch 

18from torch import nn 

19from torchvision.transforms import v2 

20 

21 

22def make_ssl_augment(img_size: int, min_scale: float = 0.4) -> v2.Transform: 

23 """Build a standard joint-embedding augmentation pipeline. 

24 

25 Args: 

26 img_size: Output image side length (square). 

27 min_scale: Lower bound of the random-resized-crop area fraction. 

28 

29 Returns: 

30 A composed transform mapping a float image batch ``[B, C, H, W]`` in ``[0, 1]`` to an 

31 augmented batch of the same shape. 

32 """ 

33 return v2.Compose( 

34 [ 

35 v2.RandomResizedCrop(size=img_size, scale=(min_scale, 1.0), antialias=True), 

36 v2.RandomHorizontalFlip(p=0.5), 

37 v2.RandomApply([v2.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8), 

38 v2.RandomGrayscale(p=0.2), 

39 ] 

40 ) 

41 

42 

43def make_light_augment(img_size: int, min_scale: float = 0.6) -> v2.Transform: 

44 """Build a light spatial augmentation pipeline (for MAE). 

45 

46 Args: 

47 img_size: Output image side length (square). 

48 min_scale: Lower bound of the random-resized-crop area fraction. 

49 

50 Returns: 

51 A composed transform mapping a float image batch to an augmented batch. 

52 """ 

53 return v2.Compose( 

54 [ 

55 v2.RandomResizedCrop(size=img_size, scale=(min_scale, 1.0), antialias=True), 

56 v2.RandomHorizontalFlip(p=0.5), 

57 ] 

58 ) 

59 

60 

61class TwoView(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env) 

62 """Produce two independently augmented views of an image batch.""" 

63 

64 def __init__(self, augment: v2.Transform) -> None: 

65 """Wrap an augmentation transform to emit a positive-pair. 

66 

67 Args: 

68 augment: The per-view augmentation to apply twice (independently). 

69 """ 

70 super().__init__() 

71 self.augment = augment 

72 

73 def forward(self, imgs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 

74 """Return two independently augmented views of ``imgs``. 

75 

76 Args: 

77 imgs: Image batch ``[B, C, H, W]``. 

78 

79 Returns: 

80 A ``(view_1, view_2)`` pair, each ``[B, C, H, W]``. 

81 """ 

82 return self.augment(imgs), self.augment(imgs)