Added models and supporting files

.gitignore

@@ -0,0 +1 @@
+__pycache__/

config.yml

@@ -0,0 +1,50 @@
+seed: 42
+
+feature_extraction_configs:
+  blink_detection: true
+  extraction_library: "mediapipe"
+  show_features: ['full_imgs', 'faces', 'eyes', 'blinks', 'iris']
+
+model_configs:
+  models_path: "pre_trained_models"
+  registered_model_names: ["ResNet18", "ResNet50"]
+  labels: ["left_eye", "right_eye"]
+  targets: ["left_pupil", "right_pupil"]
+  num_classes: 1
+
+xai_configs:
+  attribution_methods: [
+    "IntegratedGradients",
+    "Saliency",
+    "InputXGradient",
+    "GuidedBackprop",
+    "Deconvolution",
+    "GuidedGradCam",
+    "LayerGradCam",
+    "LayerGradientXActivation",
+  ]
+  cam_methods: [
+    "CAM",
+    "GradCAM",
+    "GradCAMpp",
+    "SmoothGradCAMpp",
+    "ScoreCAM",
+    "SSCAM",
+    "ISCAM",
+    "XGradCAM",
+    "LayerCAM",
+  ]
+
+use_sr: false
+
+upscale_configs:
+  upscale: [1, 2, 3, 4]
+  upscale_method_configs:
+    size: [16, 32]
+    antialias: true
+    interpolation: ["bicubic"]
+      
+sr_methods: ["GFPGAN", "RealESRGAN", "SRResNet", "CodeFormer", "HAT"]
+sr_method_configs:
+  bg_upsampler_name: "realesrgan"
+  prefered_net_in_upsampler: "RRDBNet"

pre_trained_models/ResNet18/right_eye.pt

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

pre_trained_models/ResNet50/left_eye.pt

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

pre_trained_models/ResNet50/right_eye.pt

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

registrations/models.py

@@ -0,0 +1,124 @@
+import sys
+import torch.nn as nn
+import os.path as osp
+from torchvision import models
+import torch.nn.functional as F
+from registry import MODEL_REGISTRY
+
+root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
+sys.path.append(root_path)
+
+# ============================= ResNets =============================
+
+
+# @MODEL_REGISTRY.register()
+# class ResNet18(nn.Module):
+#     def __init__(self, model_args):
+#         super(ResNet18, self).__init__()
+#         self.num_classes = model_args.get("num_classes", 1)
+#         self.resnet = models.resnet18(weights=None, num_classes=self.num_classes)
+
+#     def forward(self, x, masks=None):
+#         return self.resnet(x)
+
+
+# @MODEL_REGISTRY.register()
+# class ResNet18(nn.Module):
+#     def __init__(self, model_args):
+#         super(ResNet18, self).__init__()
+#         self.num_classes = model_args.get("num_classes", 1)
+#         self.resnet = models.resnet18(weights=None, num_classes=self.num_classes)
+
+#     def forward(self, x, masks=None):
+#         # Calculate the padding dynamically based on the input size
+#         height, width = x.shape[2], x.shape[3]
+#         pad_height = max(0, (224 - height) // 2)
+#         pad_width = max(0, (224 - width) // 2)
+
+#         # Apply padding
+#         x = F.pad(
+#             x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
+#         )
+#         x = self.resnet(x)
+#         return x
+
+
+@MODEL_REGISTRY.register()
+class ResNet18(nn.Module):
+    def __init__(self, model_args):
+        super(ResNet18, self).__init__()
+        self.num_classes = model_args.get("num_classes", 1)
+        self.resnet = models.resnet18(weights=None)
+        self.regression_head = nn.Linear(1000, self.num_classes)
+
+    def forward(self, x, masks=None):
+        # Calculate the padding dynamically based on the input size
+        height, width = x.shape[2], x.shape[3]
+        pad_height = max(0, (224 - height) // 2)
+        pad_width = max(0, (224 - width) // 2)
+
+        # Apply padding
+        x = F.pad(
+            x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
+        )
+        x = self.resnet(x)
+        x = self.regression_head(x)
+        return x
+
+
+# @MODEL_REGISTRY.register()
+# class ResNet50(nn.Module):
+#     def __init__(self, model_args):
+#         super(ResNet50, self).__init__()
+#         self.num_classes = model_args.get("num_classes", 1)
+#         self.resnet = models.resnet50(weights=None, num_classes=self.num_classes)
+
+#     def forward(self, x, masks=None):
+#         return self.resnet(x)
+
+
+# @MODEL_REGISTRY.register()
+# class ResNet50(nn.Module):
+#     def __init__(self, model_args):
+#         super(ResNet50, self).__init__()
+#         self.num_classes = model_args.get("num_classes", 1)
+#         self.resnet = models.resnet50(weights=None, num_classes=self.num_classes)
+
+#     def forward(self, x, masks=None):
+#         # Calculate the padding dynamically based on the input size
+#         height, width = x.shape[2], x.shape[3]
+#         pad_height = max(0, (224 - height) // 2)
+#         pad_width = max(0, (224 - width) // 2)
+
+#         # Apply padding
+#         x = F.pad(
+#             x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
+#         )
+#         x = self.resnet(x)
+#         return x
+
+
+@MODEL_REGISTRY.register()
+class ResNet50(nn.Module):
+    def __init__(self, model_args):
+        super(ResNet50, self).__init__()
+        self.num_classes = model_args.get("num_classes", 1)
+        self.resnet = models.resnet50(weights=None)
+        self.regression_head = nn.Linear(1000, self.num_classes)
+
+    def forward(self, x, masks=None):
+        # Calculate the padding dynamically based on the input size
+        height, width = x.shape[2], x.shape[3]
+        pad_height = max(0, (224 - height) // 2)
+        pad_width = max(0, (224 - width) // 2)
+
+        # Apply padding
+        x = F.pad(
+            x, (pad_width, pad_width, pad_height, pad_height), mode="constant", value=0
+        )
+        x = self.resnet(x)
+        x = self.regression_head(x)
+        return x
+
+
+print("Registered models in MODEL_REGISTRY:", MODEL_REGISTRY.keys())

registry.py

@@ -0,0 +1,82 @@
+# Modified from: https://github.com/facebookresearch/fvcore/blob/master/fvcore/common/registry.py  # noqa: E501
+
+
+class Registry:
+    """
+    The registry that provides name -> object mapping, to support third-party
+    users' custom modules.
+
+    To create a registry (e.g. a backbone registry):
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY = Registry('BACKBONE')
+
+    To register an object:
+
+    .. code-block:: python
+
+        @BACKBONE_REGISTRY.register()
+        class MyBackbone():
+            ...
+
+    Or:
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY.register(MyBackbone)
+    """
+
+    def __init__(self, name):
+        """
+        Args:
+            name (str): the name of this registry
+        """
+        self._name = name
+        self._obj_map = {}
+
+    def _do_register(self, name, obj):
+        assert name not in self._obj_map, (
+            f"An object named '{name}' was already registered "
+            f"in '{self._name}' registry!"
+        )
+        self._obj_map[name] = obj
+
+    def register(self, obj=None):
+        """
+        Register the given object under the the name `obj.__name__`.
+        Can be used as either a decorator or not.
+        See docstring of this class for usage.
+        """
+        if obj is None:
+            # used as a decorator
+            def deco(func_or_class):
+                name = func_or_class.__name__
+                self._do_register(name, func_or_class)
+                return func_or_class
+
+            return deco
+
+        # used as a function call
+        name = obj.__name__
+        self._do_register(name, obj)
+
+    def get(self, name):
+        ret = self._obj_map.get(name)
+        if ret is None:
+            raise KeyError(
+                f"No object named '{name}' found in '{self._name}' registry!"
+            )
+        return ret
+
+    def __contains__(self, name):
+        return name in self._obj_map
+
+    def __iter__(self):
+        return iter(self._obj_map.items())
+
+    def keys(self):
+        return self._obj_map.keys()
+
+
+MODEL_REGISTRY = Registry("model")

registry_utils.py

@@ -0,0 +1,79 @@
+import os
+import importlib
+from os import path as osp
+
+
+def scandir(dir_path, suffix=None, recursive=False, full_path=False):
+    """Scan a directory to find the interested files.
+
+    Args:
+        dir_path (str): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        full_path (bool, optional): If set to True, include the dir_path.
+            Default: False.
+
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith(".") and entry.is_file():
+                if full_path:
+                    return_path = entry.path
+                else:
+                    return_path = osp.relpath(entry.path, root)
+
+                if suffix is None:
+                    yield return_path
+                elif return_path.endswith(suffix):
+                    yield return_path
+            else:
+                if recursive:
+                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
+                else:
+                    continue
+
+    return _scandir(dir_path, suffix=suffix, recursive=recursive)
+
+
+def import_registered_modules(registration_folder="registrations"):
+    """
+    Import all registered modules from the specified folder.
+
+    This function automatically scans all the files under the specified folder and imports all the required modules for registry.
+
+    Parameters:
+        registration_folder (str, optional): Path to the folder containing registration modules. Default is "registrations".
+
+    Returns:
+        list: List of imported modules.
+    """
+
+    print("\n")
+
+    registration_modules_folder = (
+        osp.dirname(osp.abspath(__file__)) + f"/{registration_folder}"
+    )
+    print("registration_modules_folder = ", registration_modules_folder)
+
+    registration_modules_file_names = [
+        osp.splitext(osp.basename(v))[0]
+        for v in scandir(dir_path=registration_modules_folder)
+    ]
+    print("registration_modules_file_names = ", registration_modules_file_names)
+
+    imported_modules = [
+        importlib.import_module(f"{registration_folder}.{file_name}")
+        for file_name in registration_modules_file_names
+    ]
+    print("imported_modules = ", imported_modules)
+    print("\n")

requirements.txt

@@ -0,0 +1,27 @@
+tqdm
+PyYAML
+numpy
+pandas
+matplotlib
+seaborn
+mlflow
+pillow
+scikit_learn
+torch
+captum
+evaluate
+# basicsr
+facexlib
+realesrgan
+opencv_python
+cmake
+dlib
+einops
+transformers
+# gfpgan
+# streamlit
+mediapipe
+imutils
+scipy
+torchvision==0.16.0
+torchcam

utils.py

@@ -0,0 +1,11 @@
+from registry import MODEL_REGISTRY
+
+
+def get_model(model_configs):
+    registered_model = MODEL_REGISTRY.get(model_configs["registered_model_name"])
+    model_configs.pop("registered_model_name")
+    if len(model_configs) > 0:
+        model = registered_model(model_configs)
+    else:
+        model = registered_model()
+    return model