Repository: GHOST-Science-Club/tree-classification-irim
Pull request #45: 43-test-write-python-tests-and-include-them-in-the-build-pipeline
Author: login: rojberr
Labels: name: enhancement
description: New feature or request
Created: 2025-03-03T10:44:06Z | Merged: 2025-03-04T20:38:35Z

**This PR does the following:**
- add automated tests
- run mypy, flake8, pytest tests parallel using different jobs
- include how to run tests locally in README.md
- add tests pipeline badge to README.md

#43 


Diff:
diff --git a/.github/workflows/pep8-lint.yml b/.github/workflows/pep8-lint.yml
deleted file mode 100644
index df34ac3d..00000000
--- a/.github/workflows/pep8-lint.yml
+++ /dev/null
@@ -1,32 +0,0 @@
-name: PEP-8 Lint Compliance Check
-
-on:
-  pull_request:
-    branches:
-      - main
-  push:
-    branches:
-      - main
-
-jobs:
-  pep8-lint:
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.x'
-
-      - name: Install flake8
-        run: |
-          python -m pip install --upgrade pip
-          pip install flake8
-
-      - name: Lint using flake8 on /src without --exit-zero
-        run: |
-          flake8 src/ --count --select=E9,F63,F7,F82 --show-source --statistics
-          flake8 src/ --count --max-line-length=200 --statistics
diff --git a/.github/workflows/test-and-coverage.yml b/.github/workflows/test-and-coverage.yml
new file mode 100644
index 00000000..aa33a668
--- /dev/null
+++ b/.github/workflows/test-and-coverage.yml
@@ -0,0 +1,92 @@
+name: Tests and Coverage Deployment
+
+on:
+    pull_request:
+
+jobs:
+#    mypy-tests:
+#        runs-on: ${{ matrix.os }}
+#        strategy:
+#            matrix:
+#                os: [ ubuntu-latest, windows-latest ]
+#                python-version: [ '3.11' ]
+#
+#        steps:
+#            -   uses: actions/checkout@v3
+#            -   name: Set up Python ${{ matrix.python-version }}
+#                uses: actions/setup-python@v4
+#                with:
+#                    python-version: ${{ matrix.python-version }}
+#            -   name: Install dependencies
+#                run: |
+#                    python -m pip install --upgrade pip
+#                    pip install tox tox-gh-actions
+#            -   name: Run Linux TOX tests
+#                if: runner.os == 'Linux'
+#                run: tox -c tox_lin.ini -e mypy
+#            -   name: Run Windows TOX tests
+#                if: runner.os == 'Windows'
+#                run: tox -c tox_win.ini -e mypy
+
+    flake8-tests:
+        runs-on: ${{ matrix.os }}
+        strategy:
+            matrix:
+                os: [ ubuntu-latest, windows-latest ]
+                python-version: [ '3.11' ]
+
+        steps:
+            -   uses: actions/checkout@v3
+            -   name: Set up Python ${{ matrix.python-version }}
+                uses: actions/setup-python@v4
+                with:
+                    python-version: ${{ matrix.python-version }}
+            -   name: Install dependencies
+                run: |
+                    python -m pip install --upgrade pip
+                    pip install tox tox-gh-actions
+            -   name: Run Linux TOX tests
+                if: runner.os == 'Linux'
+                run: tox -c tox_lin.ini -e flake8
+            -   name: Run Windows TOX tests
+                if: runner.os == 'Windows'
+                run: tox -c tox_win.ini -e flake8
+
+    pytest:
+        runs-on: ${{ matrix.os }}
+        strategy:
+            matrix:
+                os: [ ubuntu-latest, windows-latest ]
+                python-version: [ '3.11' ]
+
+        steps:
+            -   uses: actions/checkout@v3
+            -   name: Set up Python ${{ matrix.python-version }}
+                uses: actions/setup-python@v4
+                with:
+                    python-version: ${{ matrix.python-version }}
+            -   name: Install dependencies
+                run: |
+                    python -m pip install --upgrade pip
+                    pip install tox tox-gh-actions
+            -   name: Run Linux TOX tests
+                if: runner.os == 'Linux'
+                run: tox -c tox_lin.ini
+            -   name: Run Windows TOX tests
+                if: runner.os == 'Windows'
+                run: tox -c tox_win.ini
+
+#    deploy:
+#        needs: test
+#        runs-on: ubuntu-latest
+#        permissions:
+#            pages: write
+#            id-token: write
+#        steps:
+#            -   name: Upload coverage report artifact
+#                uses: actions/upload-pages-artifact@v3
+#                with:
+#                    path: './htmlcov/'
+#            -   name: Deploy to GitHub Pages
+#                uses: actions/deploy-pages@v4
+#                id: deployment
diff --git a/.gitignore b/.gitignore
index bd05d534..ce9df09d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,4 +5,7 @@ lightning_logs
 __pycache__
 .vscode
 ghost-irim
-wandb
\ No newline at end of file
+wandb
+.pytest_cache
+.coverage
+.tox
\ No newline at end of file
diff --git a/README.md b/README.md
index 06160304..49d17cc5 100644
--- a/README.md
+++ b/README.md
@@ -5,6 +5,7 @@
 [![Forks][forks-shield]][forks-url]
 [![Stargazers][stars-shield]][stars-url]
 [![Issues][issues-shield]][issues-url]
+[![Tests](https://github.com/GHOST-Science-Club/tree-classification-irim/actions/workflows/test-and-coverage.yml/badge.svg)]
 
 <!-- MARKDOWN LINKS & IMAGES -->
 
@@ -42,7 +43,7 @@
     <a href="https://github.com/GHOST-Science-Club/tree-classification-irim/tree/docu">View Demo</a>
     ·
     <a href="https://github.com/othneildrew/Best-README-Template/issues/new?labels=bug&template=bug-report---.md">Report Bug</a> -->
-    · 
+    ·
     <a href="https://github.com/GHOST-Science-Club/tree-classification-irim/issues/new?assignees=&labels=&projects=&template=feature_request.md&title=">Request Feature</a>
     ·
   </p>
@@ -64,7 +65,9 @@
 **GHOST x IRIM** is an initiative by the [GHOST](https://ghost.put.poznan.pl) student organization to develop and test
 an AI algorithm for identifying tree species from aerial photos, focusing on Polish forests.
 
-### How to run
+### How to run the project 💡
+
+! We use Python 3.11.11. If you run into any issues, please check your Python version.
 
 Clone the repo and cd into:
 
@@ -76,10 +79,10 @@ cd tree-classification-irim
 For Linux/MacOS/Unix:
 
 ```
-python -m venv .venv
+python3.11 -m venv .venv
 source .venv/bin/activate
 pip install -r unix-requirements.txt
-python src/main.py
+python3.11 src/main.py
 ```
 
 For CUDA PyTorch installation, please refer to this link: https://pytorch.org/get-started/locally/ and select your
@@ -87,10 +90,10 @@ platform and CUDA.
 For Windows with CUDA and win32 lib:
 
 ```
-python -m venv .venv
+python3.11 -m venv .venv
 .venv\Scripts\Activate.ps1
 pip install -r requirements.txt
-python src/main.py
+python3.11 src/main.py
 ```
 
 If you are a project maintainer you should have received Wandb api key.
@@ -103,6 +106,23 @@ wandb login
 
 Provide your API key when prompted.
 
+### How to run the project 💡
+
+To run the tests locally, you can use the following command:
+
+```bash
+git clone git@github.com:GHOST-Science-Club/tree-classification-irim.git
+cd tree-classification-irim
+source .venv/bin/activate
+pip install -r unix-requirements.txt
+
+# For Windows
+tox -c tox_win.ini
+# For Linux/MacOS/Unix
+tox -c tox_lin.ini
+```
+
+
 ### Objectives:
 
 - Identify tree species sensitive to fires, storms, disease, and pests.
@@ -174,7 +194,7 @@ For participant contact details, see the [Team](#the-team) section.
 We are a passionate group of students from Poznan University of Technology (PUT), ranging from first-year undergraduates
 to final-year graduate students.
 
-**Team Leader:**  
+**Team Leader:**
 Kacper Dobek ([email@example.com](mailto:email@example.com))
 
 **Team Members:**
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..daa93854
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,30 @@
+[tool.coverage.run]
+omit = [
+    "src/__init__.py",
+    "src/visualization_functions.py"
+    ]
+
+[tool.pytest.ini_options]
+pythonpath = [
+    ".",
+    "src",
+    "tests"
+]
+addopts = [
+    "-v",
+    "--cov=src",
+    "--cov-branch",
+    "--color=yes",
+    "--cov-report=html",
+]
+markers = [
+    "transforms: pre-processing with Kornia and tranformation of images",
+    "dataset_functions: downloading and managing tree species files",
+    "dataset: managing the entire dataset for the model efficiently",
+    "callbacks: managing and monitoring learning process",
+    "git_functions: managing commits to repo and wandb",
+    "model: creating model for experiments"
+]
+filterwarnings = [
+    "ignore:Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package.*",
+]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index b773e20d..c760980e 100644
Binary files a/requirements.txt and b/requirements.txt differ
diff --git a/src/dataset.py b/src/dataset.py
index 02b3ffa4..677032ca 100644
--- a/src/dataset.py
+++ b/src/dataset.py
@@ -32,17 +32,9 @@ def __getitem__(self, idx):
         with Image.open(image_path) as img:
             # Convert to numpy array
             image = np.array(img)
-            image = image[:, :, 1:]  # Removing "near-inferred" channel
-
-        # Convert numpy array to a PIL Image if needed
-        # if isinstance(image, np.ndarray):
-        #    image = Image.fromarray(image)
-
-        # Convert image to RGB if it has an alpha channel
-        # if image.mode == "RGBA":
+            image = image[:, :, 1:] if image.shape[-1] == 4 else image  # Removing "near-inferred" channel
         # We found out that PIL conversion to RGB
         # keeps the "near-inferred" channel which was not desired
-        # image = image.convert("RGB")
 
         # Apply transformations
         if self.transform:
diff --git a/src/dataset_functions.py b/src/dataset_functions.py
index a9c4950b..0173c380 100644
--- a/src/dataset_functions.py
+++ b/src/dataset_functions.py
@@ -1,9 +1,45 @@
 import zipfile
-import numpy as np
 from pathlib import Path
+
+import numpy as np
 from huggingface_hub import hf_hub_download
 
 
+def print_extracted_files(extract_dir: Path):
+    print(f"Successfully extracted to {extract_dir}")
+
+    extracted_files = Path(extract_dir).iterdir()
+    print("Extracted files:")
+    for extracted_file in list(extracted_files)[:5]:
+        print(f"- {extracted_file.stem}")
+    if len(list(extracted_files)) > 5:
+        print(f"... and {len(list(extracted_files)) - 5} more files")
+
+
+def extract_files(file_path: Path, extract_dir: Path, main_subfolders: str):
+    with zipfile.ZipFile(file_path, "r") as zip_ref:
+        # Get list of all files in zip
+        image_file_list = zip_ref.namelist()
+
+        # Extract all files, modifying their paths
+        for image_file in image_file_list:
+            # Extract file with modified path
+            source = zip_ref.read(image_file)
+
+            # I assumed we are using aeirla imagery data. However, if needed,
+            # a simple function can be written that chooses either aerial or
+            # LiDAR data
+            target_path = extract_dir / Path(image_file).relative_to(main_subfolders["aerial_imagery"])
+
+            # Create directories if they don't exist
+            target_path.parent.mkdir(parents=True, exist_ok=True)
+
+            with open(target_path, "wb") as f:
+                f.write(source)
+
+        print_extracted_files(extract_dir)
+
+
 def download_data(species_folders: dict, main_subfolders: dict, dataset_folder: Path):
     """
     Function downloads specified data from HF (PureForest dataset)
@@ -20,38 +56,10 @@ def download_data(species_folders: dict, main_subfolders: dict, dataset_folder:
         )
 
         extract_dir = dataset_folder / filename
-        extract_dir.mkdir(exist_ok=True)
+        extract_dir.mkdir(exist_ok=True, parents=True)
 
         try:
-            with zipfile.ZipFile(file_path, "r") as zip_ref:
-                # Get list of all files in zip
-                image_file_list = zip_ref.namelist()
-
-                # Extract all files, modifying their paths
-                for image_file in image_file_list:
-                    # Extract file with modified path
-                    source = zip_ref.read(image_file)
-
-                    # I assumed we are using aeirla imagery data. However, if needed,
-                    # a simple function can be written that chooses either aerial or
-                    # LiDAR data
-                    target_path = extract_dir / Path(image_file).relative_to(main_subfolders["aerial_imagery"])
-
-                    # Create directories if they don't exist
-                    target_path.parent.mkdir(parents=True, exist_ok=True)
-
-                    with open(target_path, "wb") as f:
-                        f.write(source)
-
-                print(f"Successfully extracted to {extract_dir}")
-
-                extracted_files = Path(extract_dir).iterdir()
-                print("Extracted files:")
-                for extracted_file in list(extracted_files)[:5]:
-                    print(f"- {extracted_file.stem}")
-                if len(list(extracted_files)) > 5:
-                    print(f"... and {len(list(extracted_files)) - 5} more files")
-
+            extract_files(file_path, extract_dir, main_subfolders)
         except zipfile.BadZipFile:
             print(f"Error: {filename} is not a valid zip file")
 
diff --git a/src/model.py b/src/model.py
index 3a5ba674..07c304dd 100644
--- a/src/model.py
+++ b/src/model.py
@@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
 
         # Freeze pre-trained layers
         if freeze:
diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py
new file mode 100644
index 00000000..41ce76f8
--- /dev/null
+++ b/tests/test_callbacks.py
@@ -0,0 +1,85 @@
+import pytest
+import torch
+from src.callbacks import PrintMetricsCallback
+
+
+@pytest.fixture
+def callback():
+    """Create a fresh instance of PrintMetricsCallback for each test."""
+    return PrintMetricsCallback()
+
+
+@pytest.fixture
+def mock_trainer():
+    """Create a mock trainer with predefined callback metrics."""
+    class Trainer():
+        def __init__(self):
+            self.current_epoch = 1
+            self.callback_metrics = {
+                'train_loss': torch.tensor(0.2345),
+                'train_acc': torch.tensor(0.9123),
+                'val_loss': torch.tensor(0.3456),
+                'val_acc': torch.tensor(0.8765)
+            }
+
+    trainer = Trainer()
+
+    return trainer
+
+
+@pytest.mark.callbacks
+def test_on_train_epoch_end(callback, mock_trainer):
+    callback.on_train_epoch_end(mock_trainer, None)
+    train_loss = callback.train_metrics["loss"]
+    train_acc = callback.train_metrics["acc"]
+
+    error_msg = {
+        "loss-numb": "Invalid loss metric number in training callback",
+        "acc-numb": "Invalid accuracy metric number in training callback",
+        "invalid-loss": "Invalid loss in training callback",
+        "invalid-acc": "Invalid accuracy in training callback"
+    }
+
+    assert len(train_loss) == 1, error_msg["loss-numb"]
+    assert len(train_acc) == 1, error_msg["acc-numb"]
+    assert train_loss[0] == pytest.approx(0.2345), error_msg["invalid-loss"]
+    assert train_acc[0] == pytest.approx(0.9123), error_msg["invalid-acc"]
+
+
+@pytest.mark.callbacks
+def test_on_validation_epoch_end(callback, mock_trainer):
+    callback.on_validation_epoch_end(mock_trainer, None)
+    val_loss = callback.val_metrics["loss"]
+    val_acc = callback.val_metrics["acc"]
+
+    error_msg = {
+        "loss-numb": "Invalid loss metric number in validation callback",
+        "acc-numb": "Invalid accuracy metric number in validation callback",
+        "invalid-loss": "Invalid loss in validation callback",
+        "invalid-acc": "Invalid accuracy in validation callback"
+    }
+
+    assert len(val_loss) == 1, error_msg["loss-numb"]
+    assert len(val_acc) == 1, error_msg["acc-numb"]
+    assert val_loss[0] == pytest.approx(0.3456), error_msg["invalid-loss"]
+    assert val_acc[0] == pytest.approx(0.8765), error_msg["invalid-acc"]
+
+
+@pytest.mark.callbacks
+def test_callback_prints_correctly(callback, mock_trainer, capsys):
+    callback.on_train_epoch_end(mock_trainer, None)
+    callback.on_validation_epoch_end(mock_trainer, None)
+
+    exp_out = {
+        "train-out": "Train Loss: 0.2345, Train Acc: 0.9123",
+        "val-out": "Epoch: 1, Val Loss: 0.3456, Val Acc: 0.8765"
+    }
+
+    error_msg = {
+        "train-msg": "Incorrect training callback output",
+        "val-msg": "Incorrect validation callback output"
+    }
+
+    captured = capsys.readouterr()
+    assert exp_out["train-out"] in captured.out, error_msg["train-msg"]
+    assert exp_out["val-out"] in captured.out, error_msg["val-msg"]
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
new file mode 100644
index 00000000..63a92c20
--- /dev/null
+++ b/tests/test_dataset.py
@@ -0,0 +1,175 @@
+import pytest
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from src.dataset import ForestDataset, ForestDataModule
+
+
+@pytest.fixture
+def sample_data(tmp_path):
+    """Creates data instance with two sample images."""
+    image_path1 = tmp_path / "sample_image1.jpg"
+    image1 = Image.new("RGB", (224, 224), color=(255, 0, 0))
+    image1.save(image_path1)
+
+    image_path2 = tmp_path / "sample_image2.jpg"
+    image2 = Image.new("RGB", (224, 224), color=(255, 225, 0))
+    image2.save(image_path2)
+
+    return {
+        "paths": [image_path1, image_path2],
+        "labels": [0, 1]
+    }
+
+
+@pytest.fixture
+def dataset(sample_data):
+    """Creates a dataset instance with two sample images."""
+    return ForestDataset(sample_data["paths"], sample_data["labels"])
+
+
+@pytest.fixture
+def data_module(sample_data):
+    return ForestDataModule(
+        sample_data, sample_data, sample_data, batch_size=2
+        )
+
+
+@pytest.mark.dataset
+def test_dataset_length(dataset):
+    error_msg = f"Dataset size is not as expected (size: {len(dataset)})"
+
+    assert len(dataset) == 2, error_msg
+
+
+@pytest.mark.dataset
+def test_dataset_getitem(dataset):
+    image, label = dataset[0]
+
+    error_msg = {
+        "not-tensor": f"Image is not a Tensor (type: {type(image)})",
+        "not-rgb": f"Color channels are not RGB (channels: {image.shape[0]})",
+        "label-str": f"Labels are of incorrect type (type: {type(label)})"
+    }
+
+    assert isinstance(image, torch.Tensor), error_msg["not-tensor"]
+    # Ensure it has 3 color channels (RGB)
+    assert image.shape[0] == 3, error_msg["not-rgb"]
+    assert isinstance(label, int), error_msg["label-str"]
+
+
+@pytest.mark.dataset
+def test_transforms_applied(dataset):
+    image, _ = dataset[0]
+
+    error_msg = {
+        "not-tensor": f"Image is not a Tensor (type: f{type(image)})",
+        "invalid-min": f"Invalid min in norm (min: {torch.min(image)})",
+        "invalid-max": f"Invalid max in norm (min: {torch.max(image)})"
+    }
+
+    assert isinstance(image, torch.Tensor), error_msg["not-tensor"]
+    # Normalize check
+    assert torch.min(image) >= -1.0, error_msg["invalid-min"]
+    assert torch.max(image) <= 1.0, error_msg["invalid-max"]
+
+
+@pytest.mark.dataset
+def test_missing_file_handling(tmp_path):
+    missing_image_path = tmp_path / "missing.jpg"
+    missing_data = {"paths": [missing_image_path], "labels": [1]}
+    dataset = ForestDataset(*missing_data.values())
+
+    data_module = ForestDataModule(missing_data, missing_data, missing_data)
+
+    with pytest.raises(FileNotFoundError):
+        _ = dataset[0]
+
+    with pytest.raises(FileNotFoundError):
+        data_module.setup()
+        train_loader = data_module.train_dataloader()
+        _ = next(iter(train_loader))
+
+
+@pytest.mark.dataset
+def test_multiple_samples(dataset, sample_data):
+    image1, label1 = dataset[0]
+    image2, label2 = dataset[1]
+
+    exp_label1 = sample_data['labels'][0]
+    exp_label2 = sample_data['labels'][1]
+
+    error_msg = {
+        "invalid-size": f"Unexpected dataset size (size: {len(dataset)})",
+        "not-tensor1": f"Image is not a Tensor (type: {type(image1)})",
+        "not-tensor2": f"Image is not a Tensor (type: f{type(image2)})",
+        "invalid-label1": f"Incorrect label. Expected label: {exp_label1}",
+        "invalid-label2": f"Incorrect label. Expected label: {exp_label2}",
+    }
+
+    assert len(dataset) == 2, error_msg["invalid-size"]
+    assert isinstance(image1, torch.Tensor), error_msg["not-tensor1"]
+    assert isinstance(image2, torch.Tensor), error_msg["not-tensor2"]
+    assert label1 == exp_label1, error_msg["invalid-label1"]
+    assert label2 == exp_label2, error_msg["invalid-label2"]
+
+
+@pytest.mark.dataset
+def test_setup_creates_datasets(data_module):
+    data_module.setup()
+
+    train = data_module.train_dataset
+    val = data_module.val_dataset
+    test = data_module.test_dataset
+
+    error_msg = {
+        "train-invalid": "Train dataset in data module is invalid (None)",
+        "val-invalid": "Validation dataset in data module is invalid (None)",
+        "test-invalid": "Test dataset in data module is invalid (None)",
+        "train-forest": "Train dataset is not ForestDataset",
+        "val-forest": "Validation dataset is not ForestDataset",
+        "test-forest": "Test dataset is not ForestDataset"
+    }
+
+    assert train is not None, error_msg["train-invalid"]
+    assert val is not None, error_msg["val-invalid"]
+    assert test is not None, error_msg["test-invalid"]
+    assert isinstance(train, ForestDataset), error_msg["train-forest"]
+    assert isinstance(val, ForestDataset), error_msg["val-forest"]
+    assert isinstance(test, ForestDataset), error_msg["test-forest"]
+
+
+@pytest.mark.dataset
+def test_dataloader_returns_correct_batch_size(data_module):
+    data_module.setup()
+
+    train_loader = data_module.train_dataloader()
+    val_loader = data_module.val_dataloader()
+    test_loader = data_module.test_dataloader()
+
+    images, labels = next(iter(train_loader))
+
+    error_msg = {
+        "train-loader": "Train data loader is not DataLoader",
+        "val-loader": "Validation data loader is not DataLoader",
+        "test-loader": "Test data loader is not DataLoader",
+        "size-img": f"Incorrect image batch size (current: {images.shape[0]})",
+        "size-lbl": f"Incorrect label batch size (current: {labels.shape[0]})"
+    }
+
+    assert isinstance(train_loader, DataLoader), error_msg["train-loader"]
+    assert isinstance(val_loader, DataLoader), error_msg["val-loader"]
+    assert isinstance(test_loader, DataLoader), error_msg["test-loader"]
+    # Batch size
+    assert images.shape[0] == 2, error_msg["size-img"]
+    assert labels.shape[0] == 2, error_msg["size-lbl"]
+
+
+@pytest.mark.dataset
+def test_dataloader_shuffling(data_module):
+    data_module.setup()
+    train_loader = data_module.train_dataloader()
+
+    error_msg = "Data in train data loader in not shuffled"
+
+    assert train_loader.sampler is not None, error_msg
diff --git a/tests/test_dataset_functions.py b/tests/test_dataset_functions.py
index 73501403..4c7f8a05 100644
--- a/tests/test_dataset_functions.py
+++ b/tests/test_dataset_functions.py
@@ -1,28 +1,93 @@
-from unittest.mock import patch
-from src.dataset_functions import download_data
-
-
-def test_download_data(tmp_path):
-    # Mock input
-    species_folders = {"Castanea_sativa": "data/imagery-Castanea_sativa.zip"}
-    main_subfolders = {"aerial_imagery": "imagery/"}
-    dataset_folder = tmp_path / "dataset"
-
-    # Mock hf_hub_download
-    mock_zip_path = tmp_path / "mock.zip"
-    mock_zip_path.touch()  # Create a mock file
-    with patch("huggingface_hub.hf_hub_download", return_value=mock_zip_path):
-        # Mock zipfile.ZipFile
-        with patch("zipfile.ZipFile") as mock_zip:
-            mock_zip.return_value.__enter__.return_value.namelist.return_value = [
-                "aerial_imagery/file1.tif",
-                "aerial_imagery/file2.tif",
-            ]
-            mock_zip.return_value.__enter__.return_value.read.side_effect = lambda x: b"mock data"
-
-            # Call the function
-            download_data(species_folders, main_subfolders, dataset_folder)
-
-            # Assert the expected files are created
-            assert (dataset_folder / "species1/file1.tif").exists()
-            assert (dataset_folder / "species1/file2.tif").exists()
+import huggingface_hub.errors
+import pytest
+import yaml
+from pathlib import Path
+import zipfile
+import huggingface_hub
+from src.dataset_functions import download_data, load_dataset, extract_files
+
+
+@pytest.fixture
+def dataset_folder():
+    with open("src/config.yaml", "r") as c:
+        config = yaml.safe_load(c)
+    return Path.cwd() / config["dataset"]["folder"]
+
+
+@pytest.fixture
+def species_folders(dataset_folder, monkeypatch):
+    return {
+        "Castanea_sativa": "data/imagery-Castanea_sativa.zip",
+        "Pinus_nigra": "data/imagery-Pinus_nigra.zip"
+    }
+
+
+@pytest.fixture
+def main_subfolders():
+    return {"aerial_imagery": "imagery/"}
+
+
+@pytest.fixture
+def sample_images(dataset_folder, monkeypatch):
+    exp_path1 = dataset_folder / Path(
+        "Castanea_sativa/test/TEST-Castanea_sativa-C3-17_1_42.tiff"
+        )
+
+    exp_path2 = dataset_folder / Path(
+        "Pinus_nigra/train/TRAIN-Pinus_nigra-C7-100_1_280.tiff"
+        )
+
+    return exp_path1, exp_path2
+
+
+@pytest.mark.dataset_functions
+def test_download_data(dataset_folder,
+                       species_folders,
+                       main_subfolders,
+                       sample_images
+                       ):
+    download_data(species_folders, main_subfolders, dataset_folder)
+
+    assert sample_images[0].exists(), f"File {sample_images[0]} not found"
+    assert sample_images[1].exists(), f"File {sample_images[1]} not found"
+
+
+@pytest.mark.dataset_functions
+def test_extract_file_bad_zip_error(tmp_path, dataset_folder, main_subfolders):
+    corrupt_zip_path = tmp_path / "corrupt.zip"
+    with open(corrupt_zip_path, "wb") as f:
+        f.write(b"not a real zip file")  # Writing invalid data
+
+    with pytest.raises(zipfile.BadZipFile):
+        extract_files(corrupt_zip_path, dataset_folder, main_subfolders)
+
+
+@pytest.mark.dataset_functions
+def test_hf_download_errors(dataset_folder, main_subfolders):
+    invalid_entry = {"invalid": "entry"}
+    invalid_folder = {"invalid": 12345}
+
+    with pytest.raises(huggingface_hub.errors.EntryNotFoundError):
+        download_data(invalid_entry, main_subfolders, dataset_folder)
+
+    with pytest.raises(AttributeError):
+        download_data(invalid_folder, main_subfolders, dataset_folder)
+
+
+@pytest.mark.dataset_functions
+def test_load_dataset(dataset_folder, species_folders, sample_images):
+    exp_label_map = {
+        species: idx for idx, species in enumerate(species_folders)
+        }
+
+    dataset, label_map = load_dataset(dataset_folder, species_folders)
+
+    error_msg = {
+        "img-1": f"Path {sample_images[0]} not found",
+        "img-2": f"Path {sample_images[1]} not found",
+        "wrong-label": f"Incorrect label map. Expected: {exp_label_map}"
+    }
+
+    assert sample_images[0] in dataset["test"]["paths"], error_msg["img-1"]
+    assert sample_images[1] in dataset["train"]["paths"], error_msg["img-2"]
+    assert exp_label_map == label_map, error_msg["wrong-label"]
diff --git a/tests/test_git_functions.py b/tests/test_git_functions.py
new file mode 100644
index 00000000..af70fa07
--- /dev/null
+++ b/tests/test_git_functions.py
@@ -0,0 +1,52 @@
+import subprocess
+import string
+import pytest
+from src.git_functions import get_git_branch, generate_short_hash
+
+
+@pytest.mark.git_functions
+def test_get_git_branch_env_variable(monkeypatch):
+    monkeypatch.setenv("GITHUB_HEAD_REF", "feature-branch")
+
+    assert get_git_branch() == "feature-branch", "Invalid branch name"
+
+
+@pytest.mark.git_functions
+def test_get_git_branch_git_command(monkeypatch):
+    monkeypatch.delenv("GITHUB_HEAD_REF", raising=False)
+
+    monkeypatch.setattr(subprocess, "check_output", lambda e: b"main")
+
+    assert get_git_branch() == "main", "Invalid branch name"
+
+
+@pytest.mark.git_functions
+def test_get_git_branch_git_command_failure(monkeypatch):
+    monkeypatch.delenv("GITHUB_HEAD_REF", raising=False)
+
+    def raise_(ex):
+        raise ex
+    monkeypatch.setattr(
+        subprocess,
+        "check_output",
+        lambda e: raise_(subprocess.CalledProcessError(1, "git"))
+    )
+
+    assert get_git_branch() == "unknown-branch", "Invalid branch name"
+
+
+@pytest.mark.git_functions
+def test_generate_short_hash():
+    hash_value = generate_short_hash()
+
+    allow_char = string.ascii_lowercase + string.digits
+
+    error_msg = {
+        "not-str": "Run hash is not a string",
+        "bad-length": "Run hash has incorrect length",
+        "bad-elm": "Run hash has incorrect elements"
+    }
+
+    assert isinstance(hash_value, str), error_msg["not-str"]
+    assert len(hash_value) == 6, error_msg["bad-length"]
+    assert all(char in allow_char for char in hash_value), error_msg["bad-elm"]
diff --git a/tests/test_model.py b/tests/test_model.py
new file mode 100644
index 00000000..25593acc
--- /dev/null
+++ b/tests/test_model.py
@@ -0,0 +1,211 @@
+import pytest
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from pytorch_lightning import Trainer
+from src.model import ResNetClassifier
+from src.dataset import ForestDataset
+from src.transforms import Preprocess
+
+
+@pytest.fixture
+def model(monkeypatch):
+    model = ResNetClassifier(num_classes=2, learning_rate=1e-3, freeze=True)
+
+    # running training_step() from pl.LightningModule throws
+    # an error without a trainer
+    model.trainer = Trainer()
+
+    # if no optimizer is present, it triggers an index error
+    # as model retrieves current_lr for logs
+    for _ in range(10):
+        model.trainer.optimizers.append(
+            torch.optim.Adam(model.model.fc.parameters())
+        )
+
+    # log is not tested at it is done by pytorch-lightning
+    monkeypatch.setattr(model, "log", lambda *args, **kwargs: None)
+
+    return model
+
+
+@pytest.fixture
+def sample_batch():
+    """Fixture to create a sample batch of images and labels."""
+
+    # Simulating a batch of 4 RGB images of size 224x224
+    images = torch.randn(4, 3, 224, 224)
+
+    # Random binary labels
+    labels = torch.randint(0, 2, (4,))
+
+    return images, labels
+
+
+@pytest.fixture
+def sample_data(tmp_path):
+    """Creates data instance with two sample images."""
+    image_path1 = tmp_path / "sample_image1.jpg"
+    image1 = Image.new("RGB", (224, 224), color=(255, 0, 0))
+    image1.save(image_path1)
+
+    image_path2 = tmp_path / "sample_image2.jpg"
+    image2 = Image.new("RGB", (224, 224), color=(255, 225, 0))
+    image2.save(image_path2)
+
+    return {
+        "paths": [image_path1, image_path2],
+        "labels": [0, 1]
+    }
+
+
+@pytest.fixture
+def data_loader(sample_data):
+    """Fixture to create a DataLoader for
+    testing training and validation steps."""
+    dataset = ForestDataset(
+        sample_data["paths"], sample_data["labels"], transform=Preprocess()
+    )
+    return DataLoader(dataset, batch_size=2)
+
+
+@pytest.mark.model
+def test_model_initialization(model):
+
+    error_msg = {
+        "not-init-m": "Model is not initialized",
+        "bad-loss": f"Unexpected loss function ({model.criterion})",
+        "not-init-acc": "Accuracy is not initialized"
+    }
+
+    assert model.model is not None, error_msg["not-init-m"]
+    # May need to be changed in the future for other models
+    assert isinstance(
+        model.criterion, torch.nn.CrossEntropyLoss
+    ), error_msg["bad-loss"]
+    assert model.accuracy is not None, error_msg["not-init-m"]
+
+
+@pytest.mark.model
+@pytest.mark.parametrize("num_classes", [2, 5, 10])
+def test_model_different_class_sizes(num_classes):
+    model = ResNetClassifier(num_classes=num_classes)
+
+    error_msg = f"""
+    Model classes ({model.model.fc.out_features}) do not match
+    expected number of classes: {num_classes}.
+    """
+
+    assert model.model.fc.out_features == num_classes, error_msg
+
+
+@pytest.mark.model
+def test_forward_pass(model, sample_batch):
+    images, _ = sample_batch
+    outputs = model(images)
+
+    error_msg = f"""
+    Forward pass produces wrong number of
+    samples and classes. Current: ({outputs.shape})
+    """
+
+    # Expecting 4 samples and 2 output classes
+    assert outputs.shape == (4, 2), error_msg
+
+
+@pytest.mark.model
+def test_training_step(model, sample_batch):
+    batch = sample_batch
+
+    loss = model.training_step(batch, batch_idx=0)
+
+    assert loss is not None, "Loss is invalid"
+    assert loss.item() > 0, "Loss is less than 0"
+
+
+@pytest.mark.model
+def test_validation_step(model, sample_batch):
+    batch = sample_batch
+
+    loss = model.validation_step(batch, batch_idx=0)
+
+    assert loss is not None, "Loss is invalid"
+    assert loss.item() > 0, "Loss is less than 0"
+
+
+@pytest.mark.model
+def test_test_step(model, sample_batch):
+    batch = sample_batch
+
+    loss = model.test_step(batch, batch_idx=0)
+
+    assert loss is not None, "Loss is invalid"
+    assert loss.item() > 0, "Loss is less than 0"
+
+
+@pytest.mark.model
+def test_optimizer_configuration(model):
+    optim_config = model.configure_optimizers()
+
+    scheduler = optim_config["lr_scheduler"]["scheduler"]
+
+    error_msg = {
+        "no-opt": "configure_optimizer does not output optimizer",
+        "invalid-opt": "Optimizer is invalid (None)",
+        "no-lr": "configure_optimizer does not output learning scheduler",
+        "invalid_lr": "Scheduler is invalid (None)"
+    }
+
+    assert "optimizer" in optim_config, error_msg["no-opt"]
+    assert optim_config["optimizer"] is not None, error_msg["invalid-opt"]
+    assert "lr_scheduler" in optim_config, error_msg["no-lr"]
+    assert scheduler is not None, error_msg["invalid_lr"]
+
+
+@pytest.mark.model
+def test_model_freezing(model):
+    frozen_params = [p.requires_grad for p in model.model.parameters()]
+
+    error_msg = "Pretrained weights are not frozen"
+
+    assert all(p is False for p in frozen_params[:-2]), error_msg
+
+
+@pytest.mark.model
+def test_train_dataloader(model, data_loader):
+    batch = next(iter(data_loader))
+    loss = model.training_step(batch, batch_idx=0)
+
+    assert loss.item() > 0, "Dataloader produces loss less than 0"
+
+
+@pytest.mark.model
+def test_on_after_batch_transfer_without_transform(model, sample_batch):
+    result = model.on_after_batch_transfer(sample_batch, dataloader_idx=0)
+
+    error_msg = {
+        "img": "Image has been altered by after batch transfer",
+        "lbl": "Label has been altered by after batch transfer"
+    }
+
+    assert torch.equal(result[0], sample_batch[0]), error_msg["img"]
+    assert torch.equal(result[1], sample_batch[1]), error_msg["lbl"]
+
+
+@pytest.mark.model
+def test_on_after_batch_transfer_with_transform(model, sample_batch):
+    def mock_transform(x):
+        return x * 2
+
+    model.transform = mock_transform
+    result = model.on_after_batch_transfer(sample_batch, dataloader_idx=0)
+
+    expected_x = sample_batch[0] * 2
+
+    error_msg = {
+        "img": "Image remained unchanged after batch transfer",
+        "lbl": "Label has been altered by after batch transfer"
+    }
+
+    assert torch.equal(result[0], expected_x), error_msg["img"]
+    assert torch.equal(result[1], sample_batch[1]), error_msg["lbl"]
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
new file mode 100644
index 00000000..ce2e8b41
--- /dev/null
+++ b/tests/test_transforms.py
@@ -0,0 +1,33 @@
+import pytest
+import torch
+import numpy as np
+from src.transforms import Preprocess
+from kornia.utils import image_to_tensor
+
+
+@pytest.mark.transforms
+@pytest.mark.parametrize(
+    "input_shape",
+    # Different image sizes and channels
+    [(3, 64, 64), (1, 128, 128), (3, 32, 32)],
+)
+def test_preprocess_kornia(input_shape):
+    # Create a random NumPy image (values between 0 and 255)
+    np_image = np.random.randint(
+        0, 256, size=input_shape, dtype=np.uint8
+    )
+
+    torch_input = torch.tensor(np_image)
+    preprocess = Preprocess()
+    output = preprocess(torch_input)
+    expected_output = image_to_tensor(np_image, keepdim=True).float() / 255.0
+
+    error_msg = {
+        "bad-shape": "Output shape mismatch",
+        "no-norm": "Output values not normalized correctly",
+        "bad-type": "Output dtype is not float32"
+    }
+
+    assert output.shape == expected_output.shape, error_msg["bad-shape"]
+    assert torch.all((output >= 0) & (output <= 1)), error_msg["no-norm"]
+    assert output.dtype == torch.float32, error_msg["bad-type"]
diff --git a/tox_lin.ini b/tox_lin.ini
new file mode 100644
index 00000000..4a966e39
--- /dev/null
+++ b/tox_lin.ini
@@ -0,0 +1,32 @@
+[tox]
+minversion = 4.0.0
+env_list = py{310,311}, flake8, mypy
+isolated_build = true
+
+[gh-actions]
+python =
+    3.11: py311
+
+[testenv]
+deps =
+    -r {toxinidir}/unix-requirements.txt
+    pytest
+    pytest-cov
+setenv =
+    PYTHONPATH = {toxinidir}
+install_command =
+    python -m pip install -r unix-requirements.txt
+commands =
+    pytest --verbose --color=yes --cov=src --cov-append --cov-report=html
+
+[testenv:flake8]
+basepython = python
+deps = flake8
+commands = flake8 src tests --max-line-length=200 --statistics --count --show-source
+
+[testenv:mypy]
+basepython = python
+deps =
+    -r {toxinidir}/unix-requirements.txt
+    mypy
+commands = mypy src
diff --git a/tox_win.ini b/tox_win.ini
new file mode 100644
index 00000000..8c389891
--- /dev/null
+++ b/tox_win.ini
@@ -0,0 +1,32 @@
+[tox]
+minversion = 4.0.0
+env_list = py{310,311}, flake8, mypy
+isolated_build = true
+
+[gh-actions]
+python =
+    3.11: py311
+
+[testenv]
+deps =
+    -r {toxinidir}/requirements.txt
+    pytest
+    pytest-cov
+setenv =
+    PYTHONPATH = {toxinidir}
+install_command =
+    python -m pip install -r requirements.txt
+commands =
+    pytest --verbose --color=yes --cov=src --cov-append --cov-report=html
+
+[testenv:flake8]
+basepython = python
+deps = flake8
+commands = flake8 src tests --max-line-length=200 --statistics --count --show-source
+
+[testenv:mypy]
+basepython = python
+deps =
+    -r {toxinidir}/requirements.txt
+    mypy
+commands = mypy src
diff --git a/unix-requirements.txt b/unix-requirements.txt
index 61a2a7e0..741466e8 100644
--- a/unix-requirements.txt
+++ b/unix-requirements.txt
@@ -6,15 +6,19 @@ appnope==0.1.4
 asttokens==3.0.0
 async-timeout==5.0.1
 attrs==24.2.0
+cachetools==5.5.2
 certifi==2024.8.30
+chardet==5.2.0
 charset-normalizer==3.4.0
 click==8.1.8
 colorama==0.4.6
 comm==0.2.2
 contourpy==1.3.1
+coverage==7.6.12
 cycler==0.12.1
 debugpy==1.8.9
 decorator==5.1.1
+distlib==0.3.9
 docker-pycreds==0.4.0
 exceptiongroup==1.2.2
 executing==2.1.0
@@ -31,6 +35,7 @@ ipykernel==6.29.5
 ipython==8.30.0
 jedi==0.19.2
 Jinja2==3.1.5
+joblib==1.4.2
 jupyter_client==8.6.3
 jupyter_core==5.7.2
 kiwisolver==1.4.7
@@ -61,31 +66,34 @@ pydantic==2.10.5
 pydantic_core==2.27.2
 Pygments==2.18.0
 pyparsing==3.2.0
+pyproject-api==1.9.0
 pytest==8.3.4
+pytest-cov==6.0.0
 python-dateutil==2.9.0.post0
 pytorch-lightning==2.4.0
 PyYAML==6.0.2
 pyzmq==26.2.0
 requests==2.32.3
+scikit-learn==1.6.1
+scipy==1.15.1
 sentry-sdk==2.20.0
 setproctitle==1.3.4
 six==1.17.0
 smmap==5.0.2
 stack-data==0.6.3
 sympy==1.13.1
+threadpoolctl==3.5.0
 torch==2.5.1
 torchaudio==2.5.1
 torchmetrics==1.6.0
 torchvision==0.20.1
 tornado==6.4.2
+tox==4.24.1
 tqdm==4.67.1
 traitlets==5.14.3
 typing_extensions==4.12.2
 urllib3==2.2.3
+virtualenv==20.29.2
 wandb==0.19.3
 wcwidth==0.2.13
 yarl==1.18.3
-scikit-learn==1.6.1
-joblib==1.4.2
-scipy==1.15.1
-threadpoolctl==3.5.0


Reviews:
user: login: rojberr
body: Maybe let's add to README.md a new section with how to run tests?
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T10:51:11Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T11:19:07Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T11:19:48Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T11:22:13Z
user: login: undeMalum
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-03T17:54:03Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T14:10:47Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T14:26:04Z
user: login: rojberr
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T14:58:59Z
user: login: undeMalum
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T17:34:02Z

Review comments:
diff_hunk: @@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
path: src/model.py
user: login: rojberr
body: Looks good 🎉
created_at: 2025-03-03T10:50:06Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,19 @@
+[tox]
+minversion = 4.0.0
+env_list = py{310,311}
+isolated_build = true
+
+[gh-actions]
+python =
+    3.10: py310
+    3.11: py311
+
+[testenv]
path: tox_lin.ini
user: login: rojberr
body: Let's add flake8 and mypy here
<img width="492" alt="image" src="https://github.com/user-attachments/assets/07c9816a-5d99-4eba-a310-993f6f1245fc" />

created_at: 2025-03-03T11:19:07Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,50 @@
+name: Tests and Coverage Deployment
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        python-version: ['3.10', '3.11']
path: .github/workflows/test-and-coverage.yml
user: login: rojberr
body: Let's stick to Py3.11.11 as we discussed in our weekly meetings :)
```suggestion
        python-version: ['3.11.11']
```
created_at: 2025-03-03T11:19:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,19 @@
+[tox]
+minversion = 4.0.0
+env_list = py{310,311}
+isolated_build = true
+
+[gh-actions]
+python =
+    3.10: py310
+    3.11: py311
path: tox_lin.ini
user: login: rojberr
body: ```suggestion
    3.11: py311, mypy, flake8
```
created_at: 2025-03-03T11:22:13Z
author_association: COLLABORATOR
diff_hunk: @@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
path: src/model.py
user: login: undeMalum
body: I was thinking about adding a separate README.md in the folder section. I have seen a couple of projects do that: [pytorch-lightning](https://github.com/Lightning-AI/pytorch-lightning/tree/master/tests) and [django](https://github.com/django/django/tree/main/tests)
created_at: 2025-03-03T17:54:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,19 @@
+[tox]
+minversion = 4.0.0
+env_list = py{310,311}
+isolated_build = true
+
+[gh-actions]
+python =
+    3.10: py310
+    3.11: py311
+
+[testenv]
path: tox_lin.ini
user: login: rojberr
body: Added
created_at: 2025-03-04T14:10:47Z
author_association: COLLABORATOR
diff_hunk: @@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
path: src/model.py
user: login: rojberr
body: We could move it from main README.md there :)
created_at: 2025-03-04T14:26:04Z
author_association: COLLABORATOR
diff_hunk: @@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
path: src/model.py
user: login: rojberr
body: I do not see the need for it right now. [pytorch-lightning](https://github.com/Lightning-AI/pytorch-lightning/tree/master/tests) and [django](https://github.com/django/django/tree/main/tests) are bigger project with many maintainers and much more written tests which is what resulted in creation of README for tests alone. Let's leave it in main README.md for now ;)
created_at: 2025-03-04T14:58:59Z
author_association: COLLABORATOR
diff_hunk: @@ -13,7 +13,7 @@ def __init__(self, num_classes=2, learning_rate=1e-3, transform=None, freeze=Fal
         self.transform = transform
         self.learning_rate = learning_rate
 
-        self.model = models.resnet18(pretrained=True)
+        self.model = models.resnet18(weights='DEFAULT')
path: src/model.py
user: login: undeMalum
body: True, looks good to me. 

I would only divide the instruction into two parts: one with pytest and one with tox. 

Tox takes a lot of time to build when dependencies change so we can suggest to run it before committing to make sure everything works in a separate environment.

For more frequent testing while coding I would suggest running pytest with some options.

However, I'm not sure that's necessay.
created_at: 2025-03-04T17:34:01Z
author_association: CONTRIBUTOR

PR comments:
user: login: rojberr
created_at: 2025-03-03T10:57:03Z
body: We need to fix PEP8 lint findings @undeMalum :
```
src/dataset.py:35:71: E261 at least two spaces before inline comment
src/dataset_functions.py:9:1: W293 blank line contains whitespace
src/dataset_functions.py:41:1: E302 expected 2 blank lines, found 1
1     E261 at least two spaces before inline comment
1     E302 expected 2 blank lines, found 1
1     W293 blank line contains whitespace
```
author_association: COLLABORATOR