From 2a70c67dbe2fc69c613ed89a1600403ee94f6ce3 Mon Sep 17 00:00:00 2001
From: Haidar Jomaa <130698588+haidarjomaa@users.noreply.github.com>
Date: Sun, 24 Dec 2023 20:52:59 +0200
Subject: [PATCH] Update augmentation.py

Added the necessary documentation for the test_time_augmentation function
---
 data_processing/augmentation.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/data_processing/augmentation.py b/data_processing/augmentation.py
index 558f0ee..15a26e0 100644
--- a/data_processing/augmentation.py
+++ b/data_processing/augmentation.py
@@ -13,6 +13,20 @@
 from skimage.measure import label as label_fn
 
 def test_time_augmentation(data,i):
+    """
+    Apply test-time augmentation to input data.
+
+    Args:
+    - data (torch.Tensor): Input data with shape (channels, height, width).
+    - i (int): Index specifying the augmentation type:
+        - 0: No flip
+        - 1: Flip along the last dimension
+        - 2: Flip along the second-to-last dimension
+        - 3: Flip along both dimensions
+
+    Returns:
+    - torch.Tensor: Augmented data.
+    """
     if i == 0:
         x = data
     elif i == 1: