Inference artifact saving

.gitignore

@@ -12,3 +12,7 @@ wandb
 /data
 .idea
 .editorconfig
+checkpoints/
+segmentation_outputs/
+*.ckpt
+*.onnx

dvc.yaml

@@ -18,4 +18,4 @@ stages:
         params:
             - src/config.yaml:
         deps:
-            - src/main.py
+            - src/main.py

pyproject.toml

@@ -58,6 +58,7 @@ dependencies = [
     "networkx==3.2.1",
     "numpy==2.2.4",
     "omegaconf==2.3.0",
+    "onnx>=1.20.0",
     "packaging==25.0",
     "pandas==2.3.3",
     "parso==0.8.5",
@@ -206,4 +207,4 @@ torchaudio = [
 [[tool.uv.index]]
 name = "pytorch-cu128"
 url = "https://download.pytorch.org/whl/cu128"
-explicit = true
+explicit = true

src/config.yaml

@@ -27,7 +27,7 @@ model:
   name: "fine_grained" # currently supporting resnet18, vit and inception_v3
 
 training:
-  batch_size: 32
+  batch_size: 1
   learning_rate: 0.0001
   max_epochs: 100
   freeze: true
@@ -36,6 +36,16 @@ training:
   step_size: 4
   gamma: 0.1 
 
+  augmentation:
+    enabled: true
+    horizontal_flip_p: 0.5
+    vertical_flip_p: 0.2
+    rotation_deg: 15
+    brightness: 0.15
+    contrast: 0.15
+    saturation: 0.1
+    hue: 0.02
+
   oversample:
     oversample_factor: 4
     oversample_threshold: 1000
@@ -74,4 +84,18 @@ training:
     patience: 3
     mode: "min"
 
-device: "gpu"
+inference:
+  mask_size: 224
+  export_onnx: true
+  checkpoint_path: "checkpoints/model.ckpt"
+  # wandb_artifact: "ghost-irim/ghost-irim/model-kx2co2o3:v0"
+  wandb_artifact: "ghost-irim/ghost-irim/best-model-inference-artifact-saving-vft1nd:v0"
+  wandb_project: "ghost-irim"
+  # Set to true to run inference on raw images (ONNX-style testing)
+  # Set to false to run inference on test dataset split
+  use_raw_images: false
+  # Path to directory containing raw images (relative to project root)
+  # raw_images_path: "test_images_qgis/10122025_193353"
+  # raw_images_path: "/home/adam/Documents/programming/tree-classification-irim/src/data/Abies_alba/train"
+  raw_images_path: "/home/adam/Documents/programming/tree-classification-irim/src/data/Robinia_pseudoacacia/train"
+device: "cuda"

src/dataset.py

@@ -14,6 +14,29 @@
     config = yaml.safe_load(c)
 
 
+def get_train_transform():
+    augmentation_cfg = config["training"].get("augmentation", {})
+    if not augmentation_cfg.get("enabled", True):
+        return None
+
+    return transforms.Compose(
+        [
+            transforms.Resize((224, 224)),
+            transforms.RandomHorizontalFlip(p=augmentation_cfg.get("horizontal_flip_p", 0.5)),
+            transforms.RandomVerticalFlip(p=augmentation_cfg.get("vertical_flip_p", 0.2)),
+            transforms.RandomRotation(degrees=augmentation_cfg.get("rotation_deg", 15)),
+            transforms.ColorJitter(
+                brightness=augmentation_cfg.get("brightness", 0.15),
+                contrast=augmentation_cfg.get("contrast", 0.15),
+                saturation=augmentation_cfg.get("saturation", 0.1),
+                hue=augmentation_cfg.get("hue", 0.02),
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+
 def calculate_dataloader_params(batch_size, img_size=(224, 224), image_channels=3, precision=32, ram_fraction=0.8):
     """
     Function calculates the number of workers and prefetch factor
@@ -45,7 +68,12 @@ def calculate_dataloader_params(batch_size, img_size=(224, 224), image_channels=
         prefetch_factor = min(max_batches_in_ram, 16)
         num_workers = min(floor(prefetch_factor / 2), os.cpu_count())
 
-        params = {"num_workers": num_workers, "prefetch_factor": prefetch_factor, "pin_memory": config["device"] == "gpu", "persistent_workers": True}
+        params = {
+            "num_workers": num_workers,
+            "prefetch_factor": prefetch_factor,
+            "pin_memory": str(config.get("device", "")).startswith("cuda") or config.get("device") == "gpu",
+            "persistent_workers": True,
+        }
 
     else:
         params = {
@@ -182,7 +210,8 @@ def __init__(self, train_data, val_data, test_data, dataset, dataset_args=None,
         self.params = calculate_dataloader_params(batch_size)
 
     def setup(self, stage=None):
-        self.train_dataset = self.dataset(image_paths=self.train_data["paths"], labels=self.train_data["labels"], **self.dataset_args)
+        train_transform = get_train_transform()
+        self.train_dataset = self.dataset(image_paths=self.train_data["paths"], labels=self.train_data["labels"], transform=train_transform, **self.dataset_args)
         self.val_dataset = ForestDataset(image_paths=self.val_data["paths"], labels=self.val_data["labels"])
         self.test_dataset = ForestDataset(image_paths=self.test_data["paths"], labels=self.test_data["labels"])