From dc16527f52677de5d16a77508cffdc02d9044378 Mon Sep 17 00:00:00 2001
From: Benjamin Midtvedt <benmid@student.chalmers.se>
Date: Mon, 2 Aug 2021 23:59:58 +0200
Subject: [PATCH] Implement rudimentary citation system

---
 deeptrack/backend/citations.py    | 35 +++++++++++++++++++++++++++++++
 deeptrack/backend/core.py         | 10 +++++++++
 deeptrack/models/convolutional.py |  8 ++++---
 deeptrack/models/utils.py         | 16 ++++++++++++++
 4 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 deeptrack/backend/citations.py

diff --git a/deeptrack/backend/citations.py b/deeptrack/backend/citations.py
new file mode 100644
index 000000000..baee70f35
--- /dev/null
+++ b/deeptrack/backend/citations.py
@@ -0,0 +1,35 @@
+deeptrack_bibtex = """
+@article{Midtvet2021DeepTrack,
+    author  = {Midtvedt,Benjamin  and 
+               Helgadottir,Saga  and 
+               Argun,Aykut  and 
+               Pineda,Jesús  and 
+               Midtvedt,Daniel  and 
+               Volpe,Giovanni},
+    title   = {Quantitative digital microscopy with deep learning},
+    journal = {Applied Physics Reviews},
+    volume  = {8},
+    number  = {1},
+    pages   = {011310},
+    year    = {2021},
+    doi     = {10.1063/5.0034891}
+}
+"""
+
+unet_bibtex = """
+@article{DBLP:journals/corr/RonnebergerFB15,
+  author    = {Olaf Ronneberger and
+               Philipp Fischer and
+               Thomas Brox},
+  title     = {U-Net: Convolutional Networks for Biomedical Image Segmentation},
+  journal   = {CoRR},
+  volume    = {abs/1505.04597},
+  year      = {2015},
+  url       = {http://arxiv.org/abs/1505.04597},
+  archivePrefix = {arXiv},
+  eprint    = {1505.04597},
+  timestamp = {Mon, 13 Aug 2018 16:46:52 +0200},
+  biburl    = {https://dblp.org/rec/journals/corr/RonnebergerFB15.bib},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+"""
\ No newline at end of file
diff --git a/deeptrack/backend/core.py b/deeptrack/backend/core.py
index ab2066a0f..d4bc92f76 100644
--- a/deeptrack/backend/core.py
+++ b/deeptrack/backend/core.py
@@ -4,6 +4,7 @@
 import operator
 
 from .. import utils, image
+from .citations import deeptrack_bibtex
 
 
 class DeepTrackDataObject:
@@ -134,6 +135,8 @@ class DeepTrackNode:
 
     __nonelike_default = object()
 
+    citations = {deeptrack_bibtex}
+
     def __init__(self, action=__nonelike_default, **kwargs):
         self.data = DeepTrackDataList()
         self.children = []
@@ -248,6 +251,13 @@ def recurse_dependencies(self, memory=None):
         for dependency in self.dependencies:
             yield from dependency.recurse_dependencies(memory=memory)
 
+    def get_citations(self):
+        cites = self.citations
+        for dep in self.recurse_dependencies():
+            cites.union(dep.citations)
+
+        return cites
+
     def __call__(self, replicate_index=None):
 
         if not self.is_valid(replicate_index=replicate_index):
diff --git a/deeptrack/models/convolutional.py b/deeptrack/models/convolutional.py
index 26c74396e..458c9ddd8 100644
--- a/deeptrack/models/convolutional.py
+++ b/deeptrack/models/convolutional.py
@@ -1,7 +1,8 @@
-from .utils import as_KerasModel
+from tensorflow.keras import layers, models
 
+from ..backend.citations import unet_bibtex
 from .layers import as_block
-from tensorflow.keras import layers, models
+from .utils import as_KerasModel, with_citation
 
 
 def center_crop(layer, target_layer):
@@ -124,6 +125,7 @@ def Convolutional(
 convolutional = Convolutional
 
 
+@with_citation(unet_bibtex)
 @as_KerasModel
 def UNet(
     input_shape=(None, None, 1),
@@ -218,7 +220,7 @@ def UNet(
 
         layer = upsampling_block(conv_layer_dimension)(layer)
 
-        concat_layer = center_crop(concat_layer, layer)
+        # concat_layer = center_crop(concat_layer, layer) Not currently working
 
         layer = layers.Concatenate(axis=-1)([layer, concat_layer])
 
diff --git a/deeptrack/models/utils.py b/deeptrack/models/utils.py
index 7fc610a35..91b02e0a2 100644
--- a/deeptrack/models/utils.py
+++ b/deeptrack/models/utils.py
@@ -35,6 +35,22 @@ def load_model(filepath, custom_objects=None, compile=True, options=None):
     model = KerasModel(model, compile=False)
 
 
+def with_citation(citation):
+    def wrapper(func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            res = func(*args, **kwargs)
+            assert isinstance(
+                res, features.Feature
+            ), "Wrapped model is not a deeptrack object. Did you forget @as_KerasModel?"
+            res.citations = {*res.citations, citation}
+            return res
+
+        return inner
+
+    return wrapper
+
+
 def as_KerasModel(func):
     @wraps(func)
     def inner(*args, **kwargs):