Restructure Abstract UDF classes to allow for the future addition of EVA transforms

api-docs/source/reference/udf.rst

@@ -22,9 +22,9 @@ Implementing a custom UDF
 
     UDFs packaged along with EVA are located inside the `eva/udfs/` folder
 
-2. Next, create and implement a class that derives from `PytorchAbstractUDF`.
+2. Next, create and implement a class that derives from `PytorchClassifierAbstractUDF`.
 
-    * The `PytorchAbstractUDF` is a parent class that defines and implements standard methods for model inference.
+    * The `PytorchClassifierAbstractUDF` is a parent class that defines and implements standard methods for model inference.
 
     * `_get_predictions()` - an abstract method that needs to be implemented in your child class.
 

eva/udfs/abstract_udfs.py

@@ -15,16 +15,16 @@
 from abc import ABCMeta, abstractmethod
 from typing import List
 
-import numpy as np
+from numpy.typing import ArrayLike
 import pandas as pd
 
 from eva.models.catalog.frame_info import FrameInfo
+from eva.models.catalog.properties import ColorSpace
 
 
-class AbstractClassifierUDF(metaclass=ABCMeta):
+class AbstractUDF(metaclass=ABCMeta):
     """
-    Abstract class for UDFs. All the UDFs which perform classification
-    inherit this calls.
+    Abstract class for UDFs. All the UDFs in EVA will inherit from this.
 
     Load and initialize the machine learning model in the __init__.
 
@@ -34,15 +34,15 @@ def __init__(self):
         pass
 
     @property
-    @abstractmethod
-    def input_format(self) -> FrameInfo:
-        pass
+    def name(self) -> str:
+        return str(self)
 
     @property
-    @abstractmethod
-    def name(self) -> str:
-        pass
+    def input_format(self) -> FrameInfo:
+        return FrameInfo(-1, -1, 3, ColorSpace.RGB)
+
 
+class AbstractClassifierUDF(AbstractUDF):
     @property
     @abstractmethod
     def labels(self) -> List[str]:
@@ -52,7 +52,7 @@ def labels(self) -> List[str]:
         """
 
     @abstractmethod
-    def classify(self, frames: np.ndarray) -> pd.DataFrame:
+    def classify(self, frames: ArrayLike) -> pd.DataFrame:
         """
         Takes as input a batch of frames and returns the predictions by
         applying the classification model.
@@ -67,3 +67,22 @@ def classify(self, frames: np.ndarray) -> pd.DataFrame:
 
     def __call__(self, *args, **kwargs):
         return self.classify(*args, **kwargs)
+
+
+class AbstractTransformationUDF(AbstractUDF):
+    @abstractmethod
+    def transform(self, frames: ArrayLike) -> ArrayLike:
+        """
+        Takes as input a batch of frames and transforms them
+        by applying the frame transformation model.
+
+        Arguments:
+            frames: Input batch of frames on which prediction
+            needs to be made
+
+        Returns:
+            Transformed frames
+        """
+
+    def __call__(self, *args, **kwargs):
+        return self.transform(*args, **kwargs)

eva/udfs/fastrcnn_object_detector.py

@@ -19,7 +19,7 @@
 
 from eva.models.catalog.frame_info import FrameInfo
 from eva.models.catalog.properties import ColorSpace
-from eva.udfs.pytorch_abstract_udf import PytorchAbstractUDF
+from eva.udfs.pytorch_abstract_udf import PytorchAbstractClassifierUDF
 
 try:
     from torch import Tensor
@@ -38,7 +38,7 @@
     )
 
 
-class FastRCNNObjectDetector(PytorchAbstractUDF):
+class FastRCNNObjectDetector(PytorchAbstractClassifierUDF):
     """
     Arguments:
         threshold (float): Threshold for classifier confidence score

eva/udfs/gpu_compatible.py

@@ -12,12 +12,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from abc import ABC, abstractmethod
+from abc import ABCMeta, abstractmethod
+from typing import TypeVar
 
 
-class GPUCompatible(ABC):
+GPUCompatible = TypeVar("GPUCompatible")
+
+
+class GPUCompatible(metaclass=ABCMeta):
     @abstractmethod
-    def to_device(self, device: str):
+    def to_device(self, device: str) -> GPUCompatible:
         """
         Implement this method to enable GPU for the function being executed.
         Arguments:

eva/udfs/pytorch_abstract_udf.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import List
 
 import numpy as np
@@ -23,11 +23,11 @@
 from torchvision.transforms import Compose, transforms
 
 from eva.configuration.configuration_manager import ConfigurationManager
-from eva.udfs.abstract_udfs import AbstractClassifierUDF
+from eva.udfs.abstract_udfs import AbstractClassifierUDF, AbstractTransformationUDF
 from eva.udfs.gpu_compatible import GPUCompatible
 
 
-class PytorchAbstractUDF(AbstractClassifierUDF, nn.Module, GPUCompatible, ABC):
+class PytorchAbstractClassifierUDF(AbstractClassifierUDF, nn.Module, GPUCompatible):
     """
     A pytorch based classifier. Used to make sure we make maximum
     utilization of features provided by pytorch without reinventing the wheel.
@@ -97,11 +97,9 @@ def as_numpy(self, val: Tensor) -> np.ndarray:
         """
         return val.detach().cpu().numpy()
 
-    def to_device(self, device: str):
+    def to_device(self, device: str) -> GPUCompatible:
         """
-
-        :param device:
-        :return:
+        Required to make class a member of GPUCompatible Protocol.
         """
         return self.to(torch.device("cuda:{}".format(device)))
 
@@ -112,3 +110,11 @@ def __call__(self, *args, **kwargs):
         if isinstance(frames, pd.DataFrame):
             frames = frames.transpose().values.tolist()[0]
         return nn.Module.__call__(self, frames, **kwargs)
+
+
+class PytorchAbstractTransformationUDF(AbstractTransformationUDF, Compose):
+    def __init__(self, transforms):
+        Compose.__init__(self, transforms)
+
+    def transform(self, frames: np.ndarray) -> np.ndarray:
+        return Compose(frames)

eva/udfs/ssd_object_detector.py

@@ -21,7 +21,7 @@
 
 from eva.models.catalog.frame_info import FrameInfo
 from eva.models.catalog.properties import ColorSpace
-from eva.udfs.pytorch_abstract_udf import PytorchAbstractUDF
+from eva.udfs.pytorch_abstract_udf import PytorchAbstractClassifierUDF
 
 try:
     import torch
@@ -42,7 +42,7 @@
     )
 
 
-class SSDObjectDetector(PytorchAbstractUDF):
+class SSDObjectDetector(PytorchAbstractClassifierUDF):
     @property
     def name(self) -> str:
         return "ssd"

tutorials/apps/mnist/eva_mnist_udf.py

@@ -1,13 +1,13 @@
 import pandas as pd
 from torch import Tensor
 import torch
-from eva.udfs.pytorch_abstract_udf import PytorchAbstractUDF
+from eva.udfs.pytorch_abstract_udf import PytorchAbstractClassifierUDF
 from eva.models.catalog.frame_info import FrameInfo
 from eva.models.catalog.properties import ColorSpace
 from torchvision.transforms import Compose, ToTensor, Normalize, Grayscale
 
 
-class MnistCNN(PytorchAbstractUDF):
+class MnistCNN(PytorchAbstractClassifierUDF):
 
     @property
     def name(self) -> str: