Initial commit

.github/workflows/check_file_size.yml

@@ -0,0 +1,16 @@
+name: Check file size
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  check-file-size:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check large files
+        uses: ActionsDesk/lfs-warning@v2.0
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

.github/workflows/sync_to_HF_hub.yml

@@ -0,0 +1,20 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF: ${{ secrets.HF }}
+        run: git push --force https://soumyaprabhamaiti:$HF@huggingface.co/spaces/soumyaprabhamaiti/pet-image-segmentation-pytorch main

.github/workflows/sync_to_core_lib_github.yml

@@ -0,0 +1,21 @@
+name: Sync to Core Library GitHub Repository
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-core-lib-gh:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to core lib repo
+        env:
+          CORE_LIB_GH_PAT: ${{ secrets.CORE_LIB_GH_PAT }}
+          CORE_LIB_REPO_NAME: pet_seg_core
+        run: rm -rf .git && cd $CORE_LIB_REPO_NAME && git config --global user.email "73134224+soumya-prabha-maiti@users.noreply.github.com" && git config --global user.name "Github Actions (on behalf of Soumya Prabha Maiti)" && git init --initial-branch=main && git add . && git commit -m "Modify core library" && git remote add origin https://$CORE_LIB_GH_PAT@github.com/soumya-prabha-maiti/$CORE_LIB_REPO_NAME.git && git push -u origin main --force

.gitignore

@@ -0,0 +1,170 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.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/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# 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/
+
+data/
+logs/
+lightning_logs/
+results/
+*.ipynb
+flagged/
+*.ckpt

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Soumya Prabha Maiti
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,11 @@
+---
+title: Pet Image Segmentation using PyTorch
+emoji: 🌖
+colorFrom: blue
+colorTo: pink
+sdk: gradio
+sdk_version: 5.4.0
+app_file: run_webapp.py
+pinned: false
+license: mit
+---

pet_seg_core/.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.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/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# 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/

pet_seg_core/config.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+@dataclass
+class PetSegTrainConfig:
+    EPOCHS = 5
+    BATCH_SIZE = 8
+    FAST_DEV_RUN = False
+    TOTAL_SAMPLES = 100
+    LEARNING_RATE = 1e-3
+    TRAIN_VAL_TEST_DATA_PATH = "./data/train_val_test"
+    DEPTHWISE_SEP = False
+    CHANNELS_LIST = [16, 32, 64, 128, 256]
+    DESCRIPTION_TEXT = None
+
+
+@dataclass
+class PetSegWebappConfig:
+    MODEL_WEIGHTS_GDRIVE_FILE_ID = os.environ.get("MODEL_WEIGHTS_GDRIVE_FILE_ID")
+    MODEL_WEIGHTS_LOCAL_PATH = os.environ.get(
+        "MODEL_WEIGHTS_LOCAL_PATH", "pet-segmentation-pytorch_epoch=4-step=1840.ckpt"
+    )
+    DOWNLOAD_MODEL_WEIGTHS_FROM_GDRIVE = (
+        os.environ.get("DOWNLOAD_MODEL_WEIGTHS_FROM_GDRIVE", "True") == "True"
+    )

pet_seg_core/data.py

@@ -0,0 +1,71 @@
+import torch
+import torchvision
+from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from torchvision import transforms as T
+
+from pet_seg_core.config import PetSegTrainConfig
+
+# Define the transforms
+transform = T.Compose(
+    [
+        T.ToTensor(),
+        T.Resize((128, 128), interpolation=T.InterpolationMode.NEAREST),
+    ]
+)
+
+print(f"Downloading data")
+
+# Download the dataset
+train_val_ds = torchvision.datasets.OxfordIIITPet(
+    root=PetSegTrainConfig.TRAIN_VAL_TEST_DATA_PATH,
+    split="trainval",
+    target_types="segmentation",
+    transform=transform,
+    target_transform=transform,
+    download=True,
+)
+
+print(f"Downloaded data")
+
+# Randomly sample some samples
+if PetSegTrainConfig.TOTAL_SAMPLES > 0:
+    train_val_ds = torch.utils.data.Subset(
+        train_val_ds, torch.randperm(len(train_val_ds))[:PetSegTrainConfig.TOTAL_SAMPLES]
+    )
+
+# Split the dataset into train val and test
+train_ds, val_ds = torch.utils.data.random_split(
+    train_val_ds,
+    [int(0.8 * len(train_val_ds)), len(train_val_ds) - int(0.8 * len(train_val_ds))],
+)
+
+test_ds, val_ds = torch.utils.data.random_split(
+    val_ds,
+    [int(0.5 * len(val_ds)), len(val_ds) - int(0.5 * len(val_ds))],
+)
+
+train_dataloader = DataLoader(
+    train_ds,  # The training samples.
+    sampler=RandomSampler(train_ds),  # Select batches randomly
+    batch_size=PetSegTrainConfig.BATCH_SIZE,  # Trains with this batch size.
+    num_workers=3,
+    persistent_workers=True,
+)
+
+# For validation the order doesn't matter, so we'll just read them sequentially.
+val_dataloader = DataLoader(
+    val_ds,  # The validation samples.
+    sampler=SequentialSampler(val_ds),  # Pull out batches sequentially.
+    batch_size=PetSegTrainConfig.BATCH_SIZE,  # Evaluate with this batch size.
+    num_workers=3,
+    persistent_workers=True,
+)
+
+# For validation the order doesn't matter, so we'll just read them sequentially.
+test_dataloader = DataLoader(
+            test_ds, # The validation samples.
+            sampler = SequentialSampler(test_ds), # Pull out batches sequentially.
+            batch_size = PetSegTrainConfig.BATCH_SIZE, # Evaluate with this batch size.
+            num_workers=3,
+            persistent_workers=True,
+        )

pet_seg_core/gdrive_utils.py

@@ -0,0 +1,124 @@
+import datetime
+import io
+import json
+import os
+
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
+
+
+class GDriveUtils:
+    LOG_EVENTS = True
+
+    @staticmethod
+    def get_gdrive_service(creds_stringified: str | None = None):
+        SCOPES = ["https://www.googleapis.com/auth/drive"]
+        if not creds_stringified:
+            print(
+                "Attempting to use google drive creds from environment variable"
+            ) if GDriveUtils.LOG_EVENTS else None
+            creds_stringified = os.getenv("GOOGLE_SERVICE_ACC_CREDS")
+        creds_dict = json.loads(creds_stringified)
+        creds = service_account.Credentials.from_service_account_info(
+            creds_dict, scopes=SCOPES
+        )
+        return build("drive", "v3", credentials=creds)
+
+    @staticmethod
+    def upload_file_to_gdrive(
+        local_file_path,
+        drive_parent_folder_id: str,
+        drive_filename: str | None = None,
+        creds_stringified: str | None = None,
+    ) -> str:
+        service = GDriveUtils.get_gdrive_service(creds_stringified)
+
+        if not drive_filename:
+            drive_filename = os.path.basename(local_file_path)
+
+        file_metadata = {
+            "name": drive_filename,
+            "parents": [drive_parent_folder_id],
+        }
+        file = (
+            service.files()
+            .create(body=file_metadata, media_body=local_file_path)
+            .execute()
+        )
+        print(
+            "File uploaded, drive file id: ", file.get("id")
+        ) if GDriveUtils.LOG_EVENTS else None
+        return file.get("id")
+
+    @staticmethod
+    def upload_file_to_gdrive_sanity_check(
+        drive_parent_folder_id: str,
+        creds_stringified: str | None = None,
+    ):
+        try:
+            curr_time_utc = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            file_name = f"gdrive_upload_test_{curr_time_utc}_UTC.txt"
+            print(
+                "Creating local file to upload: ", file_name
+            ) if GDriveUtils.LOG_EVENTS else None
+            with open(file_name, "w") as f:
+                f.write(f"gdrive_upload_test_{curr_time_utc}_UTC")
+            return GDriveUtils.upload_file_to_gdrive(
+                file_name, drive_parent_folder_id, creds_stringified=creds_stringified
+            )
+        except Exception as e:
+            raise e
+        finally:
+            if os.path.exists(file_name):
+                print(
+                    "Deleting local file: ", file_name
+                ) if GDriveUtils.LOG_EVENTS else None
+                os.remove(file_name)
+
+    @staticmethod
+    def download_file_from_gdrive(
+        drive_file_id: str,
+        local_file_path: str | None = None,
+        creds_stringified: str | None = None,
+    ):
+        service = GDriveUtils.get_gdrive_service(creds_stringified)
+
+        drive_filename = service.files().get(fileId=drive_file_id, fields="name").execute().get('name')
+
+        if not local_file_path:
+            local_file_path = f"{drive_file_id}_{drive_filename}"
+
+        request = service.files().get_media(fileId=drive_file_id)
+        file = io.BytesIO()
+        downloader = MediaIoBaseDownload(file, request, chunksize= 25 * 1024 * 1024)
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+            print(f"Downloading gdrive file {drive_filename} to local file {local_file_path}: {int(status.progress() * 100)}%.") if GDriveUtils.LOG_EVENTS else None
+
+        if os.path.dirname(local_file_path):
+            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+        with open(local_file_path, "wb") as f:
+            f.write(file.getvalue())
+        print(
+            "Downloaded file locally to: ", local_file_path
+        ) if GDriveUtils.LOG_EVENTS else None
+
+    @staticmethod
+    def download_file_from_gdrive_sanity_check(
+        drive_parent_folder_id: str,
+        creds_stringified: str | None = None,
+    ):
+        file_id = GDriveUtils.upload_file_to_gdrive_sanity_check(
+            drive_parent_folder_id, creds_stringified
+        )
+        GDriveUtils.download_file_from_gdrive(
+            file_id, creds_stringified=creds_stringified
+        )
+
+    def stringify_json_creds(json_file: str, txt_file: str) -> str:
+        with open(json_file, "r") as f:
+            creds_dict = json.load(f)
+        with open(txt_file, "w") as f:
+            f.write(json.dumps(creds_dict))

pet_seg_core/model.py

@@ -0,0 +1,201 @@
+import lightning as pl
+import torch
+import torchvision.transforms.functional as TF
+from torch import nn
+from torchmetrics.functional.segmentation import mean_iou
+from torchmetrics.classification import MulticlassConfusionMatrix
+from pet_seg_core.config import PetSegTrainConfig
+from functools import partial
+
+
+class DoubleConvOriginal(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConvOriginal, self).__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class DoubleConvDepthwiseSep(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConvDepthwiseSep, self).__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(
+                in_channels,
+                in_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=in_channels,
+                bias=False,
+            ),
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=out_channels,
+                bias=False,
+            ),
+            nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class UNet(pl.LightningModule):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        channels_list=[64, 128, 256, 512],
+        depthwise_sep=False,
+    ):
+        super(UNet, self).__init__()
+        self.save_hyperparameters() 
+        
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        if depthwise_sep:
+            DoubleConv = DoubleConvDepthwiseSep
+        else:
+            DoubleConv = DoubleConvOriginal
+
+        # Encoder
+        for channels in channels_list:
+            self.encoder.append(DoubleConv(in_channels, channels))
+            in_channels = channels
+
+        # Decoder
+        for channels in channels_list[::-1]:
+            self.decoder.append(
+                nn.ConvTranspose2d(channels * 2, channels, kernel_size=2, stride=2)
+            )
+            self.decoder.append(DoubleConv(channels * 2, channels))
+
+        self.bottleneck = DoubleConv(channels_list[-1], channels_list[-1] * 2)
+        self.out = nn.Conv2d(channels_list[0], out_channels, kernel_size=1)
+
+        self.loss_fn = nn.CrossEntropyLoss()
+
+        self.iou = partial(mean_iou, num_classes=out_channels)
+        self.conf_mat = MulticlassConfusionMatrix(num_classes=out_channels)
+
+    def forward(self, x):
+        skip_connections = []
+        for i, enc_block in enumerate(self.encoder):
+            x = enc_block(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+
+        skip_connections = skip_connections[::-1]
+
+        for i in range(0, len(self.decoder), 2):
+            x = self.decoder[i](x)
+            skip_connection = skip_connections[i // 2]
+
+            if x.shape != skip_connection.shape:
+                x = TF.resize(x, size=skip_connection.shape[2:])
+
+            concat_skip = torch.cat(
+                (skip_connection, x), dim=1
+            )  # Concatenate along the channel dimension
+            x = self.decoder[i + 1](concat_skip)
+
+        x = self.out(x)
+
+        return x
+
+    def _common_step(self, batch, batch_idx, prefix):
+        x, y = batch
+        y = (y * 255 - 1).long().squeeze(1)  # move to dataloader
+        y_hat = self(x)
+
+        loss = self.loss_fn(y_hat, y)
+        self.log(f"{prefix}_loss", loss.item(), prog_bar=True)
+
+        y_hat_argmax = torch.argmax(y_hat, dim=1)
+        
+        y_hat_argmax_onehot = torch.nn.functional.one_hot(y_hat_argmax, num_classes=self.out_channels).permute(0, 3, 1, 2)
+        y_onehot = torch.nn.functional.one_hot(y, num_classes=self.out_channels).permute(0, 3, 1, 2)
+        iou = self.iou(y_hat_argmax_onehot, y_onehot)
+        # self.log(f"{prefix}_iou", iou.mean().item(), prog_bar=True)
+
+        self.conf_mat.update(y_hat_argmax, y)
+
+        return y_hat, loss
+
+    def training_step(self, batch, batch_idx):
+        y_hat, loss = self._common_step(batch, batch_idx, "train")
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        y_hat, loss = self._common_step(batch, batch_idx, "val")
+
+    def test_step(self, batch, batch_idx):
+        y_hat, loss = self._common_step(batch, batch_idx, "test")
+
+    def _common_on_epoch_end(self, prefix):
+        confmat = self.conf_mat.compute()
+        
+        for i in range(self.out_channels):
+            for j in range(self.out_channels):
+                self.log(f'{prefix}_confmat_true={i}_pred={j}', confmat[i][j].item(), prog_bar=True)
+        
+        iou = torch.zeros(self.out_channels)
+        for i in range(self.out_channels):
+            true_positive = confmat[i, i]
+            false_positive = confmat.sum(dim=0)[i] - true_positive
+            false_negative = confmat.sum(dim=1)[i] - true_positive
+            union = true_positive + false_positive + false_negative
+            if union > 0:
+                iou[i] = true_positive / union
+            else:
+                iou[i] = float('nan')
+            self.log(f'{prefix}_iou_class={i}', iou[i].item(), prog_bar=True)
+            
+        self.conf_mat.reset()
+
+    def on_train_epoch_end(self):
+        self._common_on_epoch_end("train")
+        
+    def on_validation_epoch_end(self):
+        self._common_on_epoch_end("val")
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=PetSegTrainConfig.LEARNING_RATE)

pet_seg_core/pytorch_device_utils.py

@@ -0,0 +1,14 @@
+import torch
+
+def get_best_device_pytorch():
+    device = None
+    if torch.cuda.is_available():       
+        device = torch.device("cuda")
+        print(f'{torch.cuda.device_count()} GPU(s) available. Using the GPU: {torch.cuda.get_device_name(0)}')
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+        print("Using Mac ARM64 GPU")
+    else:
+        device = torch.device("cpu")
+        print('No GPU available, using CPU')
+    return device

pet_seg_core/requirements.txt

@@ -0,0 +1,85 @@
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.3
+aiohttp==3.10.10
+aiosignal==1.3.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.4.0
+click==8.1.7
+colorama==0.4.6
+fastapi==0.115.4
+ffmpy==0.4.0
+filelock==3.16.1
+frozenlist==1.5.0
+fsspec==2024.10.0
+google-api-core==2.22.0
+google-api-python-client==2.149.0
+google-auth==2.35.0
+google-auth-httplib2==0.2.0
+googleapis-common-protos==1.65.0
+gradio==5.4.0
+gradio_client==1.4.2
+h11==0.14.0
+httpcore==1.0.6
+httplib2==0.22.0
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+Jinja2==3.1.4
+lightning==2.4.0
+lightning-utilities==0.11.8
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.1.0
+networkx==3.4.2
+numpy==1.26.4
+opencv-python==4.10.0.84
+orjson==3.10.10
+packaging==24.1
+pandas==2.2.3
+pillow==11.0.0
+propcache==0.2.0
+proto-plus==1.25.0
+protobuf==5.28.3
+pyasn1==0.6.1
+pyasn1_modules==0.4.1
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.12
+pytorch-lightning==2.4.0
+pytz==2024.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.3
+rsa==4.9
+ruff==0.7.1
+safehttpx==0.1.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.41.2
+sympy==1.13.1
+tomlkit==0.12.0
+torch==2.5.1
+torchmetrics==1.5.1
+torchvision==0.20.1
+tqdm==4.66.6
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.2
+uritemplate==4.1.1
+urllib3==2.2.3
+uvicorn==0.32.0
+websockets==12.0
+yarl==1.17.0

pet_seg_core/train.py

@@ -0,0 +1,28 @@
+import lightning as pl
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.loggers import CSVLogger
+from datetime import datetime
+import os
+
+from pet_seg_core.config import PetSegTrainConfig
+from pet_seg_core.data import train_dataloader, val_dataloader
+from pet_seg_core.model import UNet
+
+def train():
+    curr_time = datetime.now().strftime('%Y-%m-%d_%H:%M:%S.%f')
+    results_folder = f"results/{curr_time}"
+    os.mkdir(results_folder)
+    with open(f"{results_folder}/description.txt", "w") as f:
+        f.write(PetSegTrainConfig.DESCRIPTION_TEXT)
+
+    logger = CSVLogger(save_dir="", name=results_folder, version="")
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=results_folder,
+        save_top_k=-1,
+    )
+    trainer = pl.Trainer(
+        max_epochs=PetSegTrainConfig.EPOCHS, fast_dev_run=PetSegTrainConfig.FAST_DEV_RUN, logger=logger, callbacks=[checkpoint_callback], gradient_clip_val=1.0
+    )
+    model = UNet(3, 3, channels_list=PetSegTrainConfig.CHANNELS_LIST, depthwise_sep=PetSegTrainConfig.DEPTHWISE_SEP)
+
+    trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

pet_seg_core/webapp.py

@@ -0,0 +1,67 @@
+from pet_seg_core.model import UNet
+from pet_seg_core.config import PetSegWebappConfig
+from pet_seg_core.gdrive_utils import GDriveUtils
+
+from torchvision import transforms as T
+import torch
+import gradio as gr
+import numpy as np
+import cv2
+from dotenv import load_dotenv
+
+load_dotenv()
+
+device = torch.device("cpu")
+
+if PetSegWebappConfig.DOWNLOAD_MODEL_WEIGTHS_FROM_GDRIVE:
+    GDriveUtils.download_file_from_gdrive(
+        PetSegWebappConfig.MODEL_WEIGHTS_GDRIVE_FILE_ID, PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH
+    )
+
+model = UNet.load_from_checkpoint(PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH)
+model.eval()
+
+def segment_image(img):
+    img = T.ToTensor()(img).unsqueeze(0).to(device)
+    mask = model(img)
+    mask = torch.argmax(mask, dim = 1).squeeze().detach().cpu().numpy()
+    return mask
+
+def overlay_mask(img, mask, alpha=0.5):
+    # Define color mapping
+    colors = {
+        0: [255, 0, 0],   # Class 0 - Red
+        1: [0, 255, 0],   # Class 1 - Green
+        2: [0, 0, 255]    # Class 2 - Blue
+        # Add more colors for additional classes if needed
+    }
+
+    # Create a blank colored overlay image
+    overlay = np.zeros_like(img)
+
+    # Map each mask value to the corresponding color
+    for class_id, color in colors.items():
+        overlay[mask == class_id] = color
+
+    # Blend the overlay with the original image
+    output = cv2.addWeighted(img, 1 - alpha, overlay, alpha, 0)
+
+    return output
+
+def transform(img):
+    mask=segment_image(img)
+    blended_img = overlay_mask(img, mask)
+    return blended_img
+
+app = gr.Interface(
+    fn=transform, 
+    inputs=gr.Image(label="Input Image"), 
+    outputs=gr.Image(label="Image with Segmentation Overlay"), 
+    title="Image Segmentation on Pet Images",
+    description="Segment image of a pet animal into three classes: background, pet, and boundary.",
+    examples=[
+        "example_images/img1.jpg",
+        "example_images/img2.jpg",
+        "example_images/img3.jpg"
+    ]
+)

requirements.txt

@@ -0,0 +1,85 @@
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.3
+aiohttp==3.10.10
+aiosignal==1.3.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.4.0
+click==8.1.7
+colorama==0.4.6
+fastapi==0.115.4
+ffmpy==0.4.0
+filelock==3.16.1
+frozenlist==1.5.0
+fsspec==2024.10.0
+google-api-core==2.22.0
+google-api-python-client==2.149.0
+google-auth==2.35.0
+google-auth-httplib2==0.2.0
+googleapis-common-protos==1.65.0
+gradio==5.4.0
+gradio_client==1.4.2
+h11==0.14.0
+httpcore==1.0.6
+httplib2==0.22.0
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+Jinja2==3.1.4
+lightning==2.4.0
+lightning-utilities==0.11.8
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.1.0
+networkx==3.4.2
+numpy==1.26.4
+opencv-python==4.10.0.84
+orjson==3.10.10
+packaging==24.1
+pandas==2.2.3
+pillow==11.0.0
+propcache==0.2.0
+proto-plus==1.25.0
+protobuf==5.28.3
+pyasn1==0.6.1
+pyasn1_modules==0.4.1
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.12
+pytorch-lightning==2.4.0
+pytz==2024.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.3
+rsa==4.9
+ruff==0.7.1
+safehttpx==0.1.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.41.2
+sympy==1.13.1
+tomlkit==0.12.0
+torch==2.5.1
+torchmetrics==1.5.1
+torchvision==0.20.1
+tqdm==4.66.6
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.2
+uritemplate==4.1.1
+urllib3==2.2.3
+uvicorn==0.32.0
+websockets==12.0
+yarl==1.17.0

run_training.py

@@ -0,0 +1,7 @@
+if __name__ == '__main__':
+    from pet_seg_core.config import PetSegTrainConfig
+    PetSegTrainConfig.EPOCHS = 5
+    PetSegTrainConfig.TOTAL_SAMPLES = -1
+    PetSegTrainConfig.DESCRIPTION_TEXT = "UNET with RGB input and 3 channel(3 class) output, trained on all samples for 5 epochs. Will be used for webapp"
+    from pet_seg_core.train import train
+    train()

run_webapp.py

@@ -0,0 +1,2 @@
+from pet_seg_core.webapp import app
+app.launch()