feat: add code for finetuning moondream

resnet/augmentation.py

@@ -0,0 +1,51 @@
+import albumentations as a
+import numpy as np
+from albumentations.pytorch import ToTensorV2
+from hyperparams import CROP_SIZE
+
+
+preprocess_training = a.Compose(
+    [
+        a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
+        a.RandomCrop(width=CROP_SIZE, height=CROP_SIZE),
+        a.GaussNoise(),
+        a.Flip(p=0.5),
+        a.RandomRotate90(p=0.5),
+        a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ToTensorV2(),
+    ]
+)
+preprocess_validation = a.Compose(
+    [
+        a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
+        a.CenterCrop(width=CROP_SIZE, height=CROP_SIZE),
+        a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ToTensorV2(),
+    ]
+)
+
+
+def transform_training(example):
+    transformed = []
+    for pil_image in example["image"]:
+        array = np.array(pil_image.convert("RGB"))
+        # check if image is in (height, width, channel) shape
+        # if not, do a transpose
+        if array.shape[-1] != 3:
+            array = np.transpose(array, (1, 2, 0))
+        img = preprocess_training(image=array)["image"]
+        transformed.append(img)
+    example["pixel_values"] = transformed
+    return example
+
+
+def transform_validation(example):
+    transformed = []
+    for pil_image in example["image"]:
+        array = np.array(pil_image.convert("RGB"))
+        if array.shape[-1] != 3:
+            array = np.transpose(array, (1, 2, 0))
+        img = preprocess_validation(image=array)["image"]
+        transformed.append(img)
+    example["pixel_values"] = transformed
+    return example