add mean laplacian and mdct

notebooks/5.0-hfk-comparing-to-traditional.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:733ce4b03b6efa6a4a0996783c8d88dcd64e9e269b16bfcb747e81dfd78b5743
-size 11339
+oid sha256:d4e1c4becea54e3d35081c4e2802d466cd0ef6f752ce5a49eeebad3091ef7262
+size 15942

requirements.txt

@@ -8,6 +8,10 @@ torchmetrics>=0.7.0
 scikit-image
 pandas
 
+# --------- libraries for image filters ---------#
+kornia
+
+
 # --------- hydra --------- #
 hydra-core>=1.1.0
 hydra-colorlog>=1.1.0

src/models/focus_traditional.py

@@ -7,6 +7,7 @@ from pytorch_lightning import LightningModule
 from torchmetrics import MaxMetric, MeanAbsoluteError, MinMetric
 from torchmetrics.classification.accuracy import Accuracy
 import torchvision.models as models
+import kornia
 
 
 def vol4(img):
@@ -16,6 +17,31 @@ def vol4(img):
     )
 
 
+def laplacian(img):
+    img_grey = torch.mean(img, dim=0).unsqueeze(0)
+    filtered = kornia.filters.laplacian(img_grey, 3)
+    mean = torch.mean(filtered)
+    return 100 / mean  # invert mean to fit metric of lower = better
+
+
+def midfrequency_dct(img):
+    kernel = torch.tensor(
+        [
+            [
+                [1, 1, -1, -1],
+                [1, 1, -1, -1],
+                [-1, -1, 1, 1],
+                [-1, -1, 1, 1],
+            ]
+        ]
+    )
+
+    img_grey = torch.mean(img, dim=0).unsqueeze(0)
+    filtered = kornia.filters.filter2d(img_grey, kernel)
+    sum = torch.sum(filtered)
+    return 100 / sum
+
+
 class TraditionalLitModule(LightningModule):
     def __init__(
         self,
@@ -26,6 +52,7 @@ class TraditionalLitModule(LightningModule):
 
         Args:
             method (str, optional): The method to use for predicting focus. Defaults to "vol4".
+                    Possible values are: vol4, mean_laplacian, midfrequency_dct
 
         Raises:
             Exception: raises exception if method parameter is not known
@@ -34,6 +61,10 @@ class TraditionalLitModule(LightningModule):
 
         if method == "vol4":
             self.function = vol4
+        if method == "mean_laplacian":
+            self.function = laplacian
+        if method == "midfrequency_dct":
+            self.function = midfrequency_dct
 
     def forward(self, x):
         return self.function(x)