Add model blocks and tests (#7)

ioannispol · web-flow · commit 4a9882ff1628 · 2023-09-30T14:16:49.000+01:00
* add model blocks and tests

* add model and tests

* add torch to setup

* add torch to requirements

* update workflow
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -9,7 +9,9 @@
 		"dockerfile": "../Dockerfile"
 	},
 	"features": {},
-    "postCreateCommand": "echo 'Development container for UWUNet is ready!'",
+    "postCreateCommand": {
+        "post_create": ".devcontainer/setup.sh"
+    },
 
     "containerEnv": {
         "PYTHONPATH": "${containerWorkspaceFolder}"
diff --git a/.devcontainer/setup.sh b/.devcontainer/setup.sh
@@ -0,0 +1,5 @@
+#!/bin/bash
+
+pip install -e .[dev]
+pip install pytest-cov
+pre-commit install
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
@@ -14,6 +14,8 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
+    - name: Set PYTHONPATH
+      run: echo "PYTHONPATH=$PWD" >> $GITHUB_ENV
     - name: Set up Python 3.8
       uses: actions/setup-python@v2
       with:
diff --git a/.gitignore b/.gitignore
@@ -1,4 +1,160 @@
+ Byte-compiled / optimized / DLL files
 __pycache__/
-*.pyc
-*.pyo
-*.pyd
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
diff --git a/UnderWaterU-Net/__init__.py b/UnderWaterU-Net/__init__.py
diff --git a/requirements.txt b/requirements.txt
@@ -3,4 +3,5 @@ numpy
 Pillow
 tqdm
 wandb
-jupyterlab
+jupyterlab
+torch >= 2.0
diff --git a/setup.cfg b/setup.cfg
@@ -0,0 +1,17 @@
+[flake8]
+max_line_length = 120
+per-file-ignores = __init__.py:F401
+
+[autflake8]
+in-place = true
+remove-all-unused-imports = true
+
+# [tool:pytest]
+# addopts = --cov
+
+[coverage:run]
+branch = False
+
+[coverage:report]
+show_missing = True
+skip_covered = True
diff --git a/setup.py b/setup.py
@@ -1,8 +1,28 @@
-
 from setuptools import setup, find_packages
 
 setup(
-    name='UnderWaterU-Net',
-    version='0.1',
-    packages=find_packages(),
+    name='underwater_unet',
+    version='0.0.1',
+    install_requires=[
+        'importlib-metadata; python_version >= "3.9"',
+        'jupyterlab',
+        'setuptools',
+        'numpy',
+        'scipy',
+        'matplotlib',
+        'pandas',
+        'torch ~= 2.0',
+    ],
+    extras_require={
+        'dev': [
+            'pytest',
+            'pre-commit',
+            'pytest-cov',
+            'nbmake'
+        ]
+    },
+    packages=find_packages(
+        include=['underwater_unet', 'underwater_unet*'],
+        exclude=['tests', 'tests.*', 'notebooks']
+        ),
 )
diff --git a/tests/test_dummy.py b/tests/test_dummy.py
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -0,0 +1,25 @@
+import torch
+
+from underwater_unet.model import UNet
+
+
+def test_UNet():
+    # Create a UNet instance
+    n_channels = 3
+    n_classes = 2
+    model = UNet(n_channels=n_channels, n_classes=n_classes)
+
+    # Generate a random tensor of size [batch_size, channels, height, width]
+    x = torch.randn(1, n_channels, 128, 128, requires_grad=True)
+
+    # Forward pass
+    y = model(x)
+
+    # Check if the tensor has been transformed correctly
+    assert y.shape == (1, n_classes, 128, 128), f"Expected shape (1, {n_classes}, 128, 128) but got {y.shape}"
+
+    # Ensure the output is of the correct type
+    assert isinstance(y, torch.Tensor), f"Expected output type torch.Tensor but got {type(y)}"
+
+    # Ensure the values are within a reasonable range (sanity check, not strictly required)
+    assert y.max() <= 10 and y.min() >= -10, "Output values are out of a reasonable range"
diff --git a/tests/test_unet_blocks.py b/tests/test_unet_blocks.py
@@ -0,0 +1,68 @@
+import torch
+import torch.nn.functional as F
+
+from underwater_unet.unet_blocks import DoubleConv, DownConv, UpConv, OutConv
+
+
+def test_DoubleConv():
+    # Create a DoubleConv instance
+    model = DoubleConv(in_channels=3, out_channels=64)
+
+    # Generate a random tensor of size [batch_size, channels, height, width]
+    x = torch.randn(1, 3, 64, 64)
+
+    # Forward pass
+    y = model(x)
+
+    # Check the output dimensions
+    assert y.shape == (1, 64, 64, 64)
+
+    # Check if mid_channels defaults to out_channels when not provided
+    assert model.double_conv[0].out_channels == 64
+
+    # Create another instance with mid_channels specified
+    model2 = DoubleConv(in_channels=3, out_channels=64, mid_channels=32)
+
+    # Check if mid_channels is set correctly
+    assert model2.double_conv[0].out_channels == 32
+
+
+def test_DownConv():
+    model = DownConv(in_channels=3, out_channels=64)
+    x = torch.rand(1, 3, 128, 128)
+    y = model(x)
+    assert y.shape == (1, 64, 64, 64), f"Expected shape (1, 64, 64, 64) but got {y.shape}"
+    assert isinstance(y, torch.Tensor), f"Expected output type torch.Tensor but got {type(y)}"
+
+
+def test_UpConv():
+    model = UpConv(in_channels=128, out_channels=64)
+
+    x1 = torch.randn(1, 64, 32, 32)
+    x2 = torch.randn(1, 64, 64, 64)
+
+    y = model(x1, x2)
+
+    assert y.shape == (1, 64, 64, 64), f"Expected shape (1, 64, 64, 64) but got {y.shape}"
+
+    assert isinstance(y, torch.Tensor), f"Expected output type torch.Tensor but got {type(y)}"
+
+    upsampled_x1 = F.interpolate(x1, scale_factor=2, mode='bilinear', align_corners=True)
+    assert torch.allclose(upsampled_x1, model.up(x1), atol=1e-5), "Bilinear upsampling not applied correctly"
+
+
+def test_OutConv():
+    in_channels = 128
+    out_channels = 64
+    model = OutConv(in_channels=in_channels, out_channels=out_channels)
+
+    x = torch.randn(1, in_channels, 64, 64)
+
+    # Forward pass
+    y = model(x)
+
+    assert y.shape[1] == out_channels, f"Expected {out_channels} channels but got {y.shape[1]} channels"
+
+    assert y.shape[2:] == x.shape[2:], f"Expected spatial dimensions {x.shape[2:]} but got {y.shape[2:]}"
+
+    assert isinstance(y, torch.Tensor), f"Expected output type torch.Tensor but got {type(y)}"
diff --git a/underwater_unet/__init__.py b/underwater_unet/__init__.py
@@ -0,0 +1 @@
+from .model import UNet
diff --git a/underwater_unet/model.py b/underwater_unet/model.py
@@ -0,0 +1,41 @@
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+from underwater_unet.unet_blocks import (DoubleConv,
+                                         DownConv,
+                                         UpConv,
+                                         OutConv)
+
+
+class UNet(nn.Module):
+
+    def __init__(self, n_channels: int, n_classes, bilinear=False) -> None:
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+
+        self.inc = (DoubleConv(n_channels, 64))
+        self.down1 = (DownConv(64, 128))
+        self.down2 = (DownConv(128, 256))
+        self.down3 = (DownConv(256, 512))
+        factor = 2 if bilinear else 1
+        self.down4 = (DownConv(512, 1024 // factor))
+        self.up1 = (UpConv(1024, 512 // factor, bilinear))
+        self.up2 = (UpConv(512, 256 // factor, bilinear))
+        self.up3 = (UpConv(256, 128 // factor, bilinear))
+        self.up4 = (UpConv(128, 64, bilinear))
+        self.outc = (OutConv(64, n_classes))
+
+    def forward(self, x):
+        x1 = checkpoint(self.inc, x)
+        x2 = checkpoint(self.down1, x1)
+        x3 = checkpoint(self.down2, x2)
+        x4 = checkpoint(self.down3, x3)
+        x5 = checkpoint(self.down4, x4)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return logits
diff --git a/underwater_unet/unet_blocks.py b/underwater_unet/unet_blocks.py
