assignment-a1: adapts build batch

mmp/a1/main.py

@@ -14,7 +14,7 @@ def pad_to_square(img):
     return F.pad(img, padding, fill=0, padding_mode='constant')
 
 
-def build_batch(paths: Sequence[str], size=(224, 224), additional_transforms=[]) -> torch.Tensor:
+def build_batch(paths: Sequence[str], transform=None) -> torch.Tensor:
     """Exercise 1.1
 
     @param paths: A sequence (e.g. list) of strings, each specifying the location of an image file.
@@ -23,13 +23,11 @@ def build_batch(paths: Sequence[str], size=(224, 224), additional_transforms=[])
     """
     preprocess = transforms.Compose([
         transforms.Lambda(pad_to_square),
-        transforms.Resize(size)]
+        transforms.Resize((224, 224)),
-        + additional_transforms
+        *([transform] if transform is not None else []),
-        +
+        transforms.ToTensor()
-        [transforms.ToTensor(),
+    ]
-         transforms.Normalize(mean=[0.485, 0.456, 0.406],
+    )
-                              std=[0.229, 0.224, 0.225]),
-         ])
     imgs = []
 
     for path in paths:
@@ -86,13 +84,13 @@ def main():
     forward_pass(paths, batch_a, model)
 
     print("Batch B:")
-    batch_b = build_batch(paths, (400, 400))
+    batch_b = build_batch(paths, transforms.Resize((100, 100)))
     forward_pass(paths, batch_b, model)
 
     print("Batch C:")
-    batch_c = build_batch(paths, additional_transforms=[
+    batch_c = build_batch(paths, transforms.RandomVerticalFlip(1))
-                          transforms.RandomVerticalFlip(1)])
     forward_pass(paths, batch_c, model)
 
+
 if __name__ == "__main__":
     main()