add src

.gitignore

@@ -0,0 +1,148 @@
+# Created by .ignore support plugin (hsz.mobi)
+
+deep-learning-models
+.pytest_cache/
+backup
+examples-local
+
+### Python template
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# 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/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
+### JetBrains template
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff:
+.idea
+.idea/**/workspace.xml
+.idea/**/tasks.xml
+.idea/dictionaries
+
+# Sensitive or high-churn files:
+.idea/**/dataSources/
+.idea/**/dataSources.ids
+.idea/**/dataSources.xml
+.idea/**/dataSources.local.xml
+.idea/**/sqlDataSources.xml
+.idea/**/dynamic.xml
+.idea/**/uiDesigner.xml
+
+# Gradle:
+.idea/**/gradle.xml
+.idea/**/libraries
+
+# Mongo Explorer plugin:
+.idea/**/mongoSettings.xml
+
+## File-based project format:
+*.iws
+
+## Plugin-specific files:
+
+# IntelliJ
+/out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+

LICENCE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 NIH/DIR
+
+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

@@ -12,4 +12,44 @@ license: mit
 short_description: Framework for Classifying patient-based AMD in CFP images
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DeepSeeNet PyTorch
+
+This repository is a PyTorch reimplementation of the original DeepSeeNet model:
+
+https://github.com/ncbi-nlp/DeepSeeNet
+
+DeepSeeNet predicts patient-level AREDS Simplified Severity Scale scores for age-related macular degeneration (AMD) from bilateral color fundus photographs. The model follows the original DeepSeeNet design by first predicting eye-level AMD risk factors, then combining predictions from both eyes into a patient-level simplified severity score.
+
+## Tasks
+
+The implementation trains three image-level subnetworks:
+
+| Task | Classes | Output |
+|---|---:|---|
+| `ADVAMD` | 2 | late AMD absent / present |
+| `DRUS` | 3 | small/none, medium, large drusen |
+| `PIG` | 2 | pigmentary abnormality absent / present |
+
+The final AREDS simplified score is computed from bilateral predictions:
+
+- score `5` if late AMD is predicted in either eye
+- otherwise, score is based on large drusen and pigmentary abnormalities across both eyes
+- bilateral medium drusen contributes one point
+
+## Citation
+
+If you use this repository, please cite the original DeepSeeNet paper:
+
+```bibtex
+@article{peng2019deepseenet,
+  title={DeepSeeNet: A Deep Learning Model for Automated Classification of Patient-based Age-related Macular Degeneration Severity from Color Fundus Photographs},
+  author={Peng, Yifan and Dharssi, Shazia and Chen, Qingyu and Keenan, Tiarnan D. and Agr\'{o}n, Elvira and Wong, Wai T. and Chew, Emily Y. and Lu, Zhiyong},
+  journal={Ophthalmology},
+  volume={126},
+  number={4},
+  pages={565--575},
+  year={2019},
+  publisher={Elsevier},
+  doi={10.1016/j.ophtha.2018.11.015}
+}
+```

RESULTS.md

@@ -0,0 +1,87 @@
+# Checkpoint Results
+
+## ADVAMD
+
+```text
+Task: ADVAMD | endpoint: late_amd | positive_class=1
+
+Metrics
+-------
+overall_accuracy     0.9658 (0.9628-0.9689)
+sensitivity          0.8417 (0.8255-0.8576)
+specificity          0.9852 (0.9831-0.9874)
+kappa                0.8498 (0.8367-0.8632)
+auc                  0.9811 (0.9777-0.9844)
+
+Classifier metrics
+------------------
+loss                 0.1119
+exact_accuracy       0.9658 (0.9628-0.9689)
+exact_kappa          0.8498 (0.8367-0.8632)
+
+Confusion matrix (rows=true, cols=pred):
+[[11217   168]
+ [  282  1499]]
+
+Binary confusion matrix (rows=true, cols=pred):
+[[11217   168]
+ [  282  1499]]
+```
+
+## DRUS
+
+```text
+Task: DRUS | endpoint: large_drusen | positive_class=2
+
+Metrics
+-------
+overall_accuracy     0.8816 (0.8763-0.8869)
+sensitivity          0.7708 (0.7588-0.7832)
+specificity          0.9368 (0.9319-0.9418)
+kappa                0.7263 (0.7144-0.7386)
+auc                  0.9489 (0.9452-0.9524)
+
+Classifier metrics
+------------------
+loss                 0.5903
+exact_accuracy       0.7471 (0.7400-0.7542)
+exact_kappa          0.6170 (0.6066-0.6280)
+macro_ovr_auc        0.8960 (0.8919-0.9001)
+
+Confusion matrix (rows=true, cols=pred):
+[[4205  820  115]
+ [ 951 2255  440]
+ [ 182  822 3376]]
+
+Binary confusion matrix (rows=true, cols=pred):
+[[8231  555]
+ [1004 3376]]
+```
+
+## PIG
+
+```text
+Task: PIG | endpoint: pigmentary_abnormality | positive_class=1
+
+Metrics
+-------
+overall_accuracy     0.8925 (0.8874-0.8976)
+sensitivity          0.8606 (0.8502-0.8701)
+specificity          0.9113 (0.9053-0.9171)
+kappa                0.7702 (0.7594-0.7811)
+auc                  0.9498 (0.9460-0.9536)
+
+Classifier metrics
+------------------
+loss                 0.2734
+exact_accuracy       0.8925 (0.8874-0.8976)
+exact_kappa          0.7702 (0.7594-0.7811)
+
+Confusion matrix (rows=true, cols=pred):
+[[7541  734]
+ [ 682 4209]]
+
+Binary confusion matrix (rows=true, cols=pred):
+[[7541  734]
+ [ 682 4209]]
+```

app.py

@@ -0,0 +1,206 @@
+import argparse
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from augmentations import get_val_transforms
+from model import DeepSeeNet
+
+
+N_CLASSES = {
+    "ADVAMD": 2,
+    "DRUS": 3,
+    "PIG": 2,
+}
+
+LABELS = {
+    "ADVAMD": ["no_late_amd", "late_amd"],
+    "DRUS": ["small_none", "medium", "large"],
+    "PIG": ["no_pigment", "pigment"],
+}
+
+
+class AlbumentationsTransform:
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, image):
+        return self.transform(image=np.asarray(image))["image"]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint-folder", default="./checkpoints")
+    parser.add_argument("--backbone", default="inception_v3")
+    parser.add_argument("--image-size", type=int, default=1024)
+    parser.add_argument("--server-name", default="127.0.0.1")
+    parser.add_argument("--server-port", type=int, default=7860)
+    parser.add_argument("--share", action="store_true")
+    return parser.parse_args()
+
+
+def load_model(path, task, backbone, device):
+    checkpoint = torch.load(path, map_location=device)
+    checkpoint_args = checkpoint.get("args", {})
+
+    model = DeepSeeNet(
+        n_classes=N_CLASSES[task],
+        backbone=checkpoint_args.get("backbone", backbone),
+        pretrained=False,
+    ).to(device)
+
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    return model
+
+
+def load_image(image, transform, device):
+    if image is None:
+        raise ValueError("Please upload both left and right images.")
+
+    image = image.convert("RGB")
+    return transform(image).unsqueeze(0).to(device)
+
+
+@torch.no_grad()
+def predict(model, image, task):
+    logits = model(image)[0].detach().cpu()
+    probs = F.softmax(logits, dim=0)
+    pred = int(torch.argmax(logits).item())
+
+    return {
+        "prediction": pred,
+        "label": LABELS[task][pred],
+        "confidence": float(probs[pred]),
+        "probabilities": {
+            LABELS[task][i]: float(probs[i])
+            for i in range(len(LABELS[task]))
+        },
+    }
+
+
+def simplified_score(scores):
+    if scores["ADVAMD"]["left"]["prediction"] == 1 or scores["ADVAMD"]["right"]["prediction"] == 1:
+        return 5
+
+    score = 0
+    score += scores["PIG"]["left"]["prediction"] == 1
+    score += scores["PIG"]["right"]["prediction"] == 1
+    score += scores["DRUS"]["left"]["prediction"] == 2
+    score += scores["DRUS"]["right"]["prediction"] == 2
+    score += (
+        scores["DRUS"]["left"]["prediction"] == 1
+        and scores["DRUS"]["right"]["prediction"] == 1
+    )
+
+    return int(min(score, 5))
+
+
+def format_probs(probabilities):
+    return " | ".join(
+        f"{label}: {prob:.3f}"
+        for label, prob in probabilities.items()
+    )
+
+
+def model_info(args, device):
+    return f"""
+# DeepSeeNet
+
+<div style="display: grid; grid-template-columns: repeat(4, max-content); gap: 0.75rem 2rem; align-items: center;">
+  <div><b>Model</b><br><code>{args.backbone}</code></div>
+  <div><b>Input size</b><br><code>{args.image_size} × {args.image_size}</code></div>
+  <div><b>Device</b><br><code>{device.type}</code></div>
+  <div><b>Checkpoint folder</b><br><code>{args.checkpoint_folder}</code></div>
+</div>
+"""
+
+
+def make_app(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    checkpoint_folder = Path(args.checkpoint_folder)
+    transform = AlbumentationsTransform(get_val_transforms(args.image_size))
+
+    models = {
+        "ADVAMD": load_model(checkpoint_folder / "advamd.pt", "ADVAMD", args.backbone, device),
+        "DRUS": load_model(checkpoint_folder / "drus.pt", "DRUS", args.backbone, device),
+        "PIG": load_model(checkpoint_folder / "pig.pt", "PIG", args.backbone, device),
+    }
+
+    def run(left_image, right_image):
+        left = load_image(left_image, transform, device)
+        right = load_image(right_image, transform, device)
+
+        scores = {}
+        for task, model in models.items():
+            scores[task] = {
+                "left": predict(model, left, task),
+                "right": predict(model, right, task),
+            }
+
+        score = simplified_score(scores)
+
+        summary_rows = [
+            ["AREDS simplified score", score],
+            ["Left eye", f"{scores['DRUS']['left']['label']}, {scores['PIG']['left']['label']}, {scores['ADVAMD']['left']['label']}"],
+            ["Right eye", f"{scores['DRUS']['right']['label']}, {scores['PIG']['right']['label']}, {scores['ADVAMD']['right']['label']}"],
+        ]
+
+        detail_rows = []
+        for task in ["ADVAMD", "DRUS", "PIG"]:
+            for eye in ["left", "right"]:
+                result = scores[task][eye]
+                detail_rows.append(
+                    [
+                        task,
+                        eye,
+                        result["label"],
+                        f"{result['confidence']:.3f}",
+                        format_probs(result["probabilities"]),
+                    ]
+                )
+
+        return summary_rows, detail_rows
+
+    with gr.Blocks(title="DeepSeeNet") as demo:
+        gr.Markdown(model_info(args, device))
+
+        with gr.Row():
+            left_image = gr.Image(type="pil", label="Left image")
+            right_image = gr.Image(type="pil", label="Right image")
+
+        button = gr.Button("Run")
+
+        summary = gr.Dataframe(
+            headers=["Item", "Result"],
+            label="Summary",
+        )
+        details = gr.Dataframe(
+            headers=["Task", "Eye", "Prediction", "Confidence", "Probabilities"],
+            label="Model outputs",
+        )
+
+        button.click(
+            run,
+            inputs=[left_image, right_image],
+            outputs=[summary, details],
+        )
+
+    return demo
+
+
+def main():
+    args = parse_args()
+    demo = make_app(args)
+    demo.launch(
+        server_name=args.server_name,
+        server_port=args.server_port,
+        share=args.share,
+    )
+
+
+if __name__ == "__main__":
+    main()

augmentations.py

@@ -0,0 +1,85 @@
+"""
+augmentations.py
+
+Simple camera-style augmentations for color fundus photography (CFP)
+classification.
+
+Expected input:
+    RGB NumPy image, shape (H, W, 3)
+
+Dependencies:
+    pip install albumentations opencv-python
+"""
+
+import cv2
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def get_train_transforms(
+    image_size=1024,
+    mean=IMAGENET_MEAN,
+    std=IMAGENET_STD,
+):
+    return A.Compose([
+        A.Resize(image_size, image_size),
+
+        # Geometry is safe
+        A.HorizontalFlip(p=0.5),
+
+        A.ShiftScaleRotate(
+            shift_limit=0.02,
+            scale_limit=0.03,   # slightly reduced
+            rotate_limit=5,     # slightly reduced
+            border_mode=0,
+            value=0,
+            p=0.3,
+        ),
+
+        # MUCH weaker photometric changes
+        A.RandomBrightnessContrast(
+            brightness_limit=0.08,   # ↓ from 0.15
+            contrast_limit=0.08,
+            p=0.3,
+        ),
+
+        # Remove or reduce gamma
+        A.RandomGamma(
+            gamma_limit=(95, 105),   # very mild
+            p=0.2,
+        ),
+
+        # Remove hue shift entirely (important)
+        # Hue shifts are not realistic for fundus physiology
+        # -> comment this out or reduce heavily
+        # A.HueSaturationValue(...)
+
+        # Keep mild quality perturbation
+        A.OneOf([
+            A.GaussianBlur(blur_limit=(3, 5)),
+            A.Downscale(scale_min=0.85, scale_max=0.95, interpolation=cv2.INTER_LINEAR),
+            A.ImageCompression(quality_lower=80, quality_upper=100),
+        ], p=0.15),
+
+        A.Normalize(mean=mean, std=std),
+        ToTensorV2(),
+    ])
+
+
+def get_val_transforms(
+    image_size=1024,
+    mean=IMAGENET_MEAN,
+    std=IMAGENET_STD,
+):
+    """
+    Validation/test transforms.
+    """
+    return A.Compose([
+        A.Resize(image_size, image_size),
+        A.Normalize(mean=mean, std=std),
+        ToTensorV2(),
+    ])

checkpoints/advamd.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c600e3f70526da0c4d65d5e4d55a563f9a083d44da223385d7a8572e194e191e
+size 89723328

checkpoints/drus.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efde754bd4d6be4b9afc4d0c17d016f138f52239ab36c4ce74ba6b24cef16245
+size 89697156

checkpoints/pig.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fcdafe1a939b5603b3fbf76276702a60c68b657b359a0702b4cf9ae74fea8ee
+size 89696006

dataloader.py

@@ -0,0 +1,55 @@
+"""PyTorch datasets and dataloaders for AREDS fundus images."""
+
+from pathlib import Path
+from typing import Callable, Optional, Tuple, Union
+
+import pandas as pd
+import torch
+from PIL import Image
+from torch import Tensor
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+TASKS = ("ADVAMD", "DRUS", "PIG")
+
+
+DEFAULT_TRANSFORM = transforms.Compose(
+    [
+        transforms.Resize(224),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=(0.485, 0.456, 0.406),
+            std=(0.229, 0.224, 0.225),
+        ),
+    ]
+)
+
+
+class AREDSDataset(Dataset):
+    def __init__(
+        self,
+        csv_path: Union[str, Path],
+        image_root: Union[str, Path],
+        task: str,
+        transform: Optional[Callable[[Image.Image], Tensor]] = None,
+    ) -> None:
+        task = task.upper()
+        if task not in TASKS:
+            raise ValueError(f"task must be one of {TASKS}")
+        self.image_root = Path(image_root)
+        self.task = task
+        self.transform = transform or DEFAULT_TRANSFORM
+        self.data = pd.read_csv(csv_path)
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor]:
+        row = self.data.iloc[index]
+        image_path = self.image_root / row.pathname
+        image = Image.open(image_path).convert("RGB")
+        image = self.transform(image)
+        label = torch.tensor(int(row[self.task]), dtype=torch.long)
+        return image, label

model.py

@@ -0,0 +1,61 @@
+"""DeepSeeNet model definition."""
+
+from torch import Tensor, nn
+
+try:
+    import timm
+except ImportError:  # pragma: no cover - handled when timm is absent.
+    timm = None
+
+
+class DeepSeeNet(nn.Module):
+    """DeepSeeNet risk-factor classifier in PyTorch.
+
+    Args:
+        n_classes: Number of output classes.
+        backbone: Any timm model name that supports ``num_classes=0``. The
+            default uses InceptionV3.
+        pretrained: Load ImageNet weights for the backbone.
+        dropout: Dropout probability used by the classifier head.
+        freeze_backbone: If true, keep the backbone frozen and train only the
+            classifier head.
+    """
+
+    def __init__(
+        self,
+        n_classes: int = 2,
+        backbone: str = "inception_v3",
+        pretrained: bool = True,
+        dropout: float = 0.5,
+        freeze_backbone: bool = False,
+    ) -> None:
+        super().__init__()
+        if n_classes < 1:
+            raise ValueError("n_classes must be positive")
+        if timm is None:
+            raise ImportError("timm is required to build DeepSeeNet")
+
+        self.backbone_name = backbone
+        self.backbone = timm.create_model(
+            backbone,
+            pretrained=pretrained,
+            num_classes=0,
+            global_pool="avg",
+        )
+        in_features = self.backbone.num_features
+        self.classifier = nn.Sequential(
+            nn.Linear(in_features, 256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(256, 128),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(128, n_classes),
+        )
+
+        if freeze_backbone:
+            self.backbone.requires_grad_(False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        features = self.backbone(x)
+        return self.classifier(features)

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+timm
+albumentations
+numpy
+pandas
+scikit-learn
+tqdm
+gradio
+pillow

run_inference.py

@@ -0,0 +1,99 @@
+"""Run DeepSeeNet inference for AREDS simplified score."""
+
+import argparse
+import json
+
+import torch
+from PIL import Image
+
+from dataloader import DEFAULT_TRANSFORM
+from model import DeepSeeNet
+
+
+N_CLASSES = {
+    "ADVAMD": 2,
+    "DRUS": 3,
+    "PIG": 2,
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--left-image", required=True)
+    parser.add_argument("--right-image", required=True)
+    parser.add_argument("--advamd-checkpoint", required=True)
+    parser.add_argument("--drus-checkpoint", required=True)
+    parser.add_argument("--pig-checkpoint", required=True)
+    parser.add_argument("--backbone", default="inception_v3")
+    return parser.parse_args()
+
+
+def load_model(checkpoint_path: str, task: str, backbone: str, device) -> DeepSeeNet:
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    checkpoint_args = checkpoint.get("args", {})
+    model = DeepSeeNet(
+        n_classes=N_CLASSES[task],
+        backbone=checkpoint_args.get("backbone", backbone),
+        pretrained=False,
+    ).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    return model
+
+
+def load_image(path: str, device) -> torch.Tensor:
+    image = Image.open(path).convert("RGB")
+    return DEFAULT_TRANSFORM(image).unsqueeze(0).to(device)
+
+
+@torch.no_grad()
+def predict(model: DeepSeeNet, image: torch.Tensor) -> int:
+    return int(model(image).argmax(dim=1).item())
+
+
+def simplified_score(scores: dict[str, tuple[int, int]]) -> int:
+    score = 0
+    if scores["ADVAMD"][0] or scores["ADVAMD"][1]:
+        return 5
+    score += scores["PIG"][0] == 1
+    score += scores["PIG"][1] == 1
+    score += scores["DRUS"][0] == 2
+    score += scores["DRUS"][1] == 2
+    score += scores["DRUS"][0] == 1 and scores["DRUS"][1] == 1
+    return min(score, 5)
+
+
+def main() -> None:
+    args = parse_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    images = {
+        "left": load_image(args.left_image, device),
+        "right": load_image(args.right_image, device),
+    }
+    checkpoints = {
+        "ADVAMD": args.advamd_checkpoint,
+        "DRUS": args.drus_checkpoint,
+        "PIG": args.pig_checkpoint,
+    }
+
+    scores = {}
+    for task, checkpoint in checkpoints.items():
+        model = load_model(checkpoint, task, args.backbone, device)
+        scores[task] = (
+            predict(model, images["left"]),
+            predict(model, images["right"]),
+        )
+
+    print(
+        json.dumps(
+            {
+                "simplified_score": simplified_score(scores),
+                "risk_factors": scores,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

test.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Callable
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+try:
+    from sklearn.metrics import (
+        accuracy_score,
+        cohen_kappa_score,
+        confusion_matrix,
+        recall_score,
+        roc_auc_score,
+    )
+except ImportError as exc:
+    raise ImportError(
+        "This evaluation script needs scikit-learn. Install with: pip install scikit-learn"
+    ) from exc
+
+from augmentations import get_val_transforms
+from dataloader import AREDSDataset
+from model import DeepSeeNet
+
+
+N_CLASSES = {
+    "ADVAMD": 2,
+    "DRUS": 3,
+    "PIG": 2,
+}
+
+DEFAULT_POSITIVE_CLASS = {
+    "ADVAMD": 1,
+    "DRUS": 2,
+    "PIG": 1,
+}
+
+ENDPOINT_NAME = {
+    "ADVAMD": "late_amd",
+    "DRUS": "large_drusen",
+    "PIG": "pigmentary_abnormality",
+}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--test-csv", required=True)
+    parser.add_argument("--image-root", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--task", required=True, type=str.upper, choices=N_CLASSES)
+    parser.add_argument("--backbone", default="inception_v3")
+    parser.add_argument("--image-size", type=int, default=1024)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=16)
+
+    parser.add_argument("--positive-class", type=int, default=None)
+    parser.add_argument("--bootstrap-iters", type=int, default=2000)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--bootstrap-unit-column", default=None)
+    parser.add_argument("--output-dir", default=None)
+    return parser.parse_args()
+
+
+class AlbumentationsTransform:
+    def __init__(self, transform) -> None:
+        self.transform = transform
+
+    def __call__(self, image):
+        return self.transform(image=np.asarray(image))["image"]
+
+
+@torch.no_grad()
+def collect_predictions(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> dict[str, np.ndarray | float]:
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+    all_labels: list[np.ndarray] = []
+    all_logits: list[np.ndarray] = []
+
+    for images, labels in tqdm(loader, desc="test"):
+        images = images.to(device)
+        labels = labels.to(device)
+
+        logits = model(images)
+        if isinstance(logits, (tuple, list)):
+            logits = logits[0]
+
+        loss = F.cross_entropy(logits, labels)
+        batch_size = labels.size(0)
+        total_loss += loss.item() * batch_size
+        total_samples += batch_size
+
+        all_labels.append(labels.detach().cpu().numpy())
+        all_logits.append(logits.detach().cpu().numpy())
+
+    labels_np = np.concatenate(all_labels).astype(int)
+    logits_np = np.concatenate(all_logits, axis=0)
+    probs_np = torch.softmax(torch.from_numpy(logits_np), dim=1).numpy()
+    preds_np = probs_np.argmax(axis=1).astype(int)
+
+    return {
+        "loss": float(total_loss / max(total_samples, 1)),
+        "labels": labels_np,
+        "logits": logits_np,
+        "probs": probs_np,
+        "preds": preds_np,
+    }
+
+
+def specificity_score(y_true_bin: np.ndarray, y_pred_bin: np.ndarray) -> float:
+    tn = np.sum((y_true_bin == 0) & (y_pred_bin == 0))
+    fp = np.sum((y_true_bin == 0) & (y_pred_bin == 1))
+    denom = tn + fp
+    return float(tn / denom) if denom else float("nan")
+
+
+def safe_auc(y_true_bin: np.ndarray, y_score: np.ndarray) -> float:
+    if len(np.unique(y_true_bin)) < 2:
+        return float("nan")
+    return float(roc_auc_score(y_true_bin, y_score))
+
+
+def compute_metrics(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    probs: np.ndarray,
+    n_classes: int,
+    positive_class: int,
+) -> dict[str, float]:
+    y_true_bin = (y_true == positive_class).astype(int)
+    y_pred_bin = (y_pred == positive_class).astype(int)
+    pos_score = probs[:, positive_class]
+
+    metrics = {
+        "loss": float("nan"),
+        "exact_accuracy": float(accuracy_score(y_true, y_pred)),
+        "exact_kappa": float(cohen_kappa_score(y_true, y_pred)),
+        "overall_accuracy": float(accuracy_score(y_true_bin, y_pred_bin)),
+        "sensitivity": float(recall_score(y_true_bin, y_pred_bin, pos_label=1, zero_division=0)),
+        "specificity": specificity_score(y_true_bin, y_pred_bin),
+        "kappa": float(cohen_kappa_score(y_true_bin, y_pred_bin)),
+        "auc": safe_auc(y_true_bin, pos_score),
+    }
+
+    if n_classes > 2 and len(np.unique(y_true)) > 1:
+        try:
+            metrics["macro_ovr_auc"] = float(
+                roc_auc_score(y_true, probs, labels=list(range(n_classes)), multi_class="ovr", average="macro")
+            )
+        except ValueError:
+            metrics["macro_ovr_auc"] = float("nan")
+
+    return metrics
+
+
+def make_bootstrap_indices(
+    n: int,
+    n_iters: int,
+    rng: np.random.Generator,
+    units: np.ndarray | None = None,
+) -> list[np.ndarray]:
+    if n_iters <= 0:
+        return []
+
+    if units is None:
+        return [rng.integers(0, n, size=n) for _ in range(n_iters)]
+
+    unique_units = np.array(pd.unique(units))
+    row_indices_by_unit = {unit: np.where(units == unit)[0] for unit in unique_units}
+    out = []
+    for _ in range(n_iters):
+        sampled_units = rng.choice(unique_units, size=len(unique_units), replace=True)
+        out.append(np.concatenate([row_indices_by_unit[u] for u in sampled_units]))
+    return out
+
+
+def bootstrap_ci(
+    metric_fn: Callable[[np.ndarray], dict[str, float]],
+    indices: list[np.ndarray],
+) -> dict[str, dict[str, float]]:
+    if not indices:
+        return {}
+
+    values_by_metric: dict[str, list[float]] = {}
+    for idx in tqdm(indices, desc="bootstrap", leave=False):
+        vals = metric_fn(idx)
+        for key, value in vals.items():
+            values_by_metric.setdefault(key, []).append(value)
+
+    intervals: dict[str, dict[str, float]] = {}
+    for key, values in values_by_metric.items():
+        arr = np.asarray(values, dtype=float)
+        intervals[key] = {
+            "ci_low": float(np.nanpercentile(arr, 2.5)),
+            "ci_high": float(np.nanpercentile(arr, 97.5)),
+        }
+    return intervals
+
+
+def combine_with_ci(metrics: dict[str, float], ci: dict[str, dict[str, float]]) -> dict[str, Any]:
+    out: dict[str, Any] = {}
+    for key, value in metrics.items():
+        out[key] = {"value": float(value)}
+        if key in ci:
+            out[key].update(ci[key])
+    return out
+
+
+def print_metric_table(metrics_with_ci: dict[str, Any]) -> None:
+    print("\nMetrics")
+    print("-------")
+    for key in ["overall_accuracy", "sensitivity", "specificity", "kappa", "auc"]:
+        item = metrics_with_ci[key]
+        if "ci_low" in item:
+            print(f"{key:20s} {item['value']:.4f} ({item['ci_low']:.4f}-{item['ci_high']:.4f})")
+        else:
+            print(f"{key:20s} {item['value']:.4f}")
+
+    print("\nClassifier metrics")
+    print("------------------")
+    for key in ["loss", "exact_accuracy", "exact_kappa", "macro_ovr_auc"]:
+        if key not in metrics_with_ci:
+            continue
+        item = metrics_with_ci[key]
+        if "ci_low" in item:
+            print(f"{key:20s} {item['value']:.4f} ({item['ci_low']:.4f}-{item['ci_high']:.4f})")
+        else:
+            print(f"{key:20s} {item['value']:.4f}")
+
+
+def main() -> None:
+    args = parse_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    task = args.task.upper()
+    n_classes = N_CLASSES[task]
+    positive_class = DEFAULT_POSITIVE_CLASS[task] if args.positive_class is None else args.positive_class
+    if not 0 <= positive_class < n_classes:
+        raise ValueError(f"positive_class={positive_class} is invalid for task={task} with {n_classes} classes")
+
+    dataset = AREDSDataset(
+        args.test_csv,
+        args.image_root,
+        task,
+        transform=AlbumentationsTransform(get_val_transforms(args.image_size)),
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+
+    model = DeepSeeNet(
+        n_classes=n_classes,
+        backbone=args.backbone,
+        pretrained=False,
+    ).to(device)
+    checkpoint = torch.load(args.checkpoint, map_location=device)
+    model.load_state_dict(checkpoint["model"])
+
+    pred_dict = collect_predictions(model, loader, device)
+    y_true = pred_dict["labels"]
+    y_pred = pred_dict["preds"]
+    probs = pred_dict["probs"]
+
+    metrics = compute_metrics(y_true, y_pred, probs, n_classes=n_classes, positive_class=positive_class)
+    metrics["loss"] = float(pred_dict["loss"])
+
+    units = None
+    if args.bootstrap_unit_column:
+        df_for_units = pd.read_csv(args.test_csv)
+        if args.bootstrap_unit_column not in df_for_units.columns:
+            raise KeyError(
+                f"--bootstrap-unit-column {args.bootstrap_unit_column!r} not found in {args.test_csv}. "
+                f"Available columns: {list(df_for_units.columns)}"
+            )
+        if len(df_for_units) != len(y_true):
+            raise ValueError(
+                "CSV length does not match dataset length. "
+                f"CSV rows={len(df_for_units)}, dataset rows={len(y_true)}"
+            )
+        units = df_for_units[args.bootstrap_unit_column].to_numpy()
+
+    rng = np.random.default_rng(args.seed)
+    bs_indices = make_bootstrap_indices(
+        n=len(y_true),
+        n_iters=args.bootstrap_iters,
+        rng=rng,
+        units=units,
+    )
+
+    def metric_fn(idx: np.ndarray) -> dict[str, float]:
+        out = compute_metrics(
+            y_true[idx],
+            y_pred[idx],
+            probs[idx],
+            n_classes=n_classes,
+            positive_class=positive_class,
+        )
+        out.pop("loss", None)
+        return out
+
+    ci = bootstrap_ci(metric_fn, bs_indices)
+    metrics_with_ci = combine_with_ci(metrics, ci)
+
+    cm = confusion_matrix(y_true, y_pred, labels=list(range(n_classes)))
+    endpoint_cm = confusion_matrix(
+        (y_true == positive_class).astype(int),
+        (y_pred == positive_class).astype(int),
+        labels=[0, 1],
+    )
+
+    meta = {
+        "task": task,
+        "endpoint": ENDPOINT_NAME[task],
+        "positive_class": int(positive_class),
+        "n_classes": int(n_classes),
+        "n_samples": int(len(y_true)),
+        "bootstrap_iters": int(args.bootstrap_iters),
+        "bootstrap_unit_column": args.bootstrap_unit_column,
+    }
+
+    print(f"\nTask: {task} | endpoint: {ENDPOINT_NAME[task]} | positive_class={positive_class}")
+    print_metric_table(metrics_with_ci)
+    print("\nConfusion matrix (rows=true, cols=pred):")
+    print(cm)
+    print("\nBinary confusion matrix (rows=true, cols=pred):")
+    print(endpoint_cm)
+
+    if args.output_dir:
+        output_dir = Path(args.output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        with (output_dir / "metrics.json").open("w") as f:
+            json.dump({"meta": meta, "metrics": metrics_with_ci}, f, indent=2)
+
+        pd.DataFrame(cm).to_csv(output_dir / "confusion_matrix.csv", index=False)
+        pd.DataFrame(endpoint_cm, index=["true_neg", "true_pos"], columns=["pred_neg", "pred_pos"]).to_csv(
+            output_dir / "endpoint_confusion_matrix.csv"
+        )
+
+        pred_df = pd.read_csv(args.test_csv)
+        if len(pred_df) == len(y_true):
+            pred_df = pred_df.copy()
+        else:
+            pred_df = pd.DataFrame(index=np.arange(len(y_true)))
+        pred_df["y_true"] = y_true
+        pred_df["y_pred"] = y_pred
+        pred_df[f"y_true_{ENDPOINT_NAME[task]}"] = (y_true == positive_class).astype(int)
+        pred_df[f"y_pred_{ENDPOINT_NAME[task]}"] = (y_pred == positive_class).astype(int)
+        for c in range(n_classes):
+            pred_df[f"prob_class_{c}"] = probs[:, c]
+        pred_df.to_csv(output_dir / "predictions.csv", index=False)
+
+        print(f"\nSaved outputs to: {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,366 @@
+import argparse
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from augmentations import get_train_transforms, get_val_transforms
+from dataloader import AREDSDataset
+from model import DeepSeeNet
+
+
+N_CLASSES = {
+    "ADVAMD": 2,
+    "DRUS": 3,
+    "PIG": 2,
+}
+
+
+class AlbumentationsTransform:
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, image):
+        return self.transform(image=np.asarray(image))["image"]
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def get_class_weights(dataset, task, device):
+    labels = torch.tensor(dataset.data[task].to_numpy(), dtype=torch.long)
+    counts = torch.bincount(labels, minlength=N_CLASSES[task]).clamp_min(1)
+    weights = counts.sum() / (len(counts) * counts)
+    return weights.to(device)
+
+
+def build_scheduler(optimizer, args):
+    if args.scheduler == "cosine":
+        return torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            T_max=args.epochs,
+            eta_min=args.min_lr,
+        )
+
+    if args.scheduler == "step":
+        return torch.optim.lr_scheduler.StepLR(
+            optimizer,
+            step_size=args.step_size,
+            gamma=args.gamma,
+        )
+
+    return None
+
+
+def train_one_epoch(
+    model,
+    loader,
+    optimizer,
+    scaler,
+    criterion,
+    device,
+    use_amp=True,
+    grad_clip=0.0,
+):
+    model.train()
+
+    running_loss = 0.0
+    running_correct = 0
+    running_samples = 0
+
+    pbar = tqdm(loader, desc="Train", leave=False)
+
+    for images, labels in pbar:
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        optimizer.zero_grad(set_to_none=True)
+
+        with torch.amp.autocast("cuda", enabled=use_amp and device.type == "cuda"):
+            logits = model(images)
+            loss = criterion(logits, labels)
+
+        if scaler is not None:
+            scaler.scale(loss).backward()
+
+            if grad_clip > 0:
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+
+            if grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+
+            optimizer.step()
+
+        batch_size = labels.size(0)
+        running_loss += loss.item() * batch_size
+        running_correct += (logits.argmax(dim=1) == labels).sum().item()
+        running_samples += batch_size
+
+        pbar.set_postfix(
+            loss=f"{running_loss / running_samples:.4f}",
+            acc=f"{running_correct / running_samples:.4f}",
+        )
+
+    return running_loss / running_samples, running_correct / running_samples
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion, device, use_amp=True):
+    model.eval()
+
+    running_loss = 0.0
+    running_correct = 0
+    running_samples = 0
+
+    pbar = tqdm(loader, desc="Val", leave=False)
+
+    for images, labels in pbar:
+        images = images.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+
+        with torch.amp.autocast("cuda", enabled=use_amp and device.type == "cuda"):
+            logits = model(images)
+            loss = criterion(logits, labels)
+
+        batch_size = labels.size(0)
+        running_loss += loss.item() * batch_size
+        running_correct += (logits.argmax(dim=1) == labels).sum().item()
+        running_samples += batch_size
+
+        pbar.set_postfix(
+            loss=f"{running_loss / running_samples:.4f}",
+            acc=f"{running_correct / running_samples:.4f}",
+        )
+
+    return running_loss / running_samples, running_correct / running_samples
+
+
+def save_checkpoint(path, model, optimizer, epoch, best_val_loss, args):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    torch.save(
+        {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "best_val_loss": best_val_loss,
+            "args": vars(args),
+        },
+        path,
+    )
+
+
+def main(args):
+    set_seed(args.seed)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    use_amp = args.amp and device.type == "cuda"
+
+    train_dataset = AREDSDataset(
+        args.train_csv,
+        args.image_root,
+        args.task,
+        transform=AlbumentationsTransform(get_train_transforms(args.image_size)),
+    )
+
+    val_dataset = AREDSDataset(
+        args.valid_csv,
+        args.image_root,
+        args.task,
+        transform=AlbumentationsTransform(get_val_transforms(args.image_size)),
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+
+    model = DeepSeeNet(
+        n_classes=N_CLASSES[args.task],
+        backbone=args.backbone,
+        pretrained=not args.no_pretrained,
+        freeze_backbone=args.freeze_backbone,
+    ).to(device)
+
+    class_weights = None
+    if not args.no_class_weights:
+        class_weights = get_class_weights(train_dataset, args.task, device)
+
+    criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+    )
+
+    scheduler = build_scheduler(optimizer, args)
+    scaler = torch.amp.GradScaler("cuda") if use_amp else None
+
+    wandb = None
+    if args.wandb:
+        import wandb
+
+        wandb.init(project=args.wandb_project, config=vars(args))
+
+    output_dir = Path(args.output_dir)
+    best_val_loss = float("inf")
+
+    print(f"Device: {device}")
+    print(f"Task: {args.task}")
+    print(f"Train samples: {len(train_dataset)}")
+    print(f"Val samples: {len(val_dataset)}")
+    print(f"Image size: {args.image_size}")
+    print(f"Batch size: {args.batch_size}")
+    print(f"Pretrained: {not args.no_pretrained}")
+    if class_weights is not None:
+        print(f"Class weights: {class_weights.detach().cpu().tolist()}")
+
+    for epoch in range(1, args.epochs + 1):
+        print(f"\nEpoch [{epoch:03d}/{args.epochs}]")
+
+        train_loss, train_acc = train_one_epoch(
+            model=model,
+            loader=train_loader,
+            optimizer=optimizer,
+            scaler=scaler,
+            criterion=criterion,
+            device=device,
+            use_amp=args.amp,
+            grad_clip=args.grad_clip,
+        )
+
+        val_loss, val_acc = evaluate(
+            model=model,
+            loader=val_loader,
+            criterion=torch.nn.CrossEntropyLoss(),
+            device=device,
+            use_amp=args.amp,
+        )
+
+        lr = optimizer.param_groups[0]["lr"]
+
+        print(
+            f"train_loss={train_loss:.4f} "
+            f"train_acc={train_acc:.4f} "
+            f"val_loss={val_loss:.4f} "
+            f"val_acc={val_acc:.4f} "
+            f"lr={lr:.2e}"
+        )
+
+        if wandb is not None:
+            wandb.log(
+                {
+                    "epoch": epoch,
+                    "lr": lr,
+                    "train_loss": train_loss,
+                    "train_acc": train_acc,
+                    "val_loss": val_loss,
+                    "val_acc": val_acc,
+                }
+            )
+
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            save_checkpoint(
+                output_dir / "best.pt",
+                model,
+                optimizer,
+                epoch,
+                best_val_loss,
+                args,
+            )
+            print(f"Saved best checkpoint: val_loss={best_val_loss:.4f}")
+
+        if args.save_every > 0 and epoch % args.save_every == 0:
+            save_checkpoint(
+                output_dir / f"epoch_{epoch:03d}.pt",
+                model,
+                optimizer,
+                epoch,
+                best_val_loss,
+                args,
+            )
+
+        if scheduler is not None:
+            scheduler.step()
+
+    save_checkpoint(
+        output_dir / "last.pt",
+        model,
+        optimizer,
+        args.epochs,
+        best_val_loss,
+        args,
+    )
+
+    print("Training complete.")
+    print(f"Best val loss: {best_val_loss:.4f}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train DeepSeeNet.")
+
+    parser.add_argument("--train-csv", required=True)
+    parser.add_argument("--valid-csv", required=True)
+    parser.add_argument("--image-root", required=True)
+    parser.add_argument("--task", required=True, type=str.upper, choices=N_CLASSES)
+    parser.add_argument("--output-dir", default="checkpoints/deepseenet")
+
+    parser.add_argument("--backbone", default="inception_v3")
+    parser.add_argument("--image-size", type=int, default=1024)
+    parser.add_argument("--epochs", type=int, default=20)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=4)
+
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--weight-decay", type=float, default=1e-4)
+    parser.add_argument("--no-pretrained", action="store_true")
+    parser.add_argument("--freeze-backbone", action="store_true")
+    parser.add_argument("--no-class-weights", action="store_true")
+
+    parser.add_argument("--scheduler", choices=("none", "cosine", "step"), default="cosine")
+    parser.add_argument("--min-lr", type=float, default=1e-6)
+    parser.add_argument("--step-size", type=int, default=5)
+    parser.add_argument("--gamma", type=float, default=0.5)
+
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument("--grad-clip", type=float, default=0.0)
+    parser.add_argument("--save-every", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--wandb-project", default="deepseenet")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)