Ocr code

src/expression/function_expression.py

@@ -115,7 +115,7 @@ def _gpu_enabled_function(self):
         if isinstance(self._function, GPUCompatible):
             device = self._context.gpu_device()
             if device != NO_GPU:
-                return self._function.to_device(device)
+                return self._function.assign_device(device)
         return self._function
 
     def __eq__(self, other):

src/udfs/gpu_compatible.py

@@ -13,9 +13,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import abstractmethod, ABC
-
+from src.constants import NO_GPU
 
 class GPUCompatible(ABC):
+
+    def __init__(self):
+        self._device = NO_GPU
+
+    @property
+    def device(self,):
+        return self._device
+
+    def assign_device(self, device: str) -> object:
+        '''
+        Assigning the device passed by the UDF to the member variable of the class
+        This is an internal function used by eva for device allocation.
+        Arguments:
+            device (str): device details
+        Returns:
+            A GPU compatible object
+        '''
+        self._device = device
+        return self.to_device(device)
+
     @abstractmethod
     def to_device(self, device: str):
         """

src/udfs/optical_character_recognition.py

@@ -0,0 +1,119 @@
+
+import pandas as pd
+import numpy as np
+import easyocr
+
+
+from typing import List
+from src.models.catalog.frame_info import FrameInfo
+from src.models.catalog.properties import ColorSpace
+from src.udfs.abstract_udfs import AbstractClassifierUDF
+from src.udfs.gpu_compatible import GPUCompatible
+
+
+class OpticalCharacterRecognition(AbstractClassifierUDF, GPUCompatible):
+    """
+    Arguments:
+        threshold (float): Threshold for classifier confidence score
+
+    """
+
+    @property
+    def name(self) -> str:
+        return "opticalcharacterrecognition"
+
+    def to_device(self, device: str):
+        """
+
+        :param device:
+        :return:
+        """
+        self.model = easyocr.Reader(['en'], gpu = "cuda:{}".format(device))
+        return self
+
+
+    def __init__(self, threshold=0.85):
+        super().__init__()
+        self.threshold = threshold
+        self.model = easyocr.Reader(['en'])
+
+
+    @property
+    def input_format(self) -> FrameInfo:
+        return FrameInfo(-1, -1, 3, ColorSpace.RGB)
+
+
+    @property
+    def labels(self) -> List[str]:
+        """
+            Empty as there are no labels required for
+            optical character recognition
+        """
+        return
+
+    def classify(self, frames: np.ndarray) -> pd.DataFrame:
+        """
+        Performs predictions on input frames
+        Arguments:
+            frames (tensor): Frames on which OCR needs
+            to be performed
+
+        Returns:
+            tuple containing predicted_classes (List[List[str]]),
+            predicted_boxes (List[List[BoundingBox]]),
+            predicted_scores (List[List[float]])
+
+        """
+
+        frames_list = frames.values.tolist()
+        frames = np.array(frames_list)
+        outcome = pd.DataFrame()
+
+        final_batch_frames = frames[0]
+        for i in range(1,frames.shape[0]):
+            final_batch_frames = np.vstack((final_batch_frames,frames[i]))
+
+        prediction = self.model.readtext_batched(final_batch_frames)
+
+        prediction = np.array(prediction)
+
+        if prediction.size != 0:
+            for detection in prediction:
+
+                pred_class = []
+                pred_boxes = []
+                pred_score = []
+                if len(detection) != 0:
+
+                    pred_class.append(detection[0][1])
+                    pred_boxes.append([[detection[0][0][0], detection[0][0][1]],
+                                [detection[0][0][2], detection[0][0][3]]])
+                    pred_score.append(detection[0][2])
+
+
+                    pred_t = \
+                        [pred_score.index(x) for x in pred_score if
+                        x > self.threshold]
+
+                    pred_t = np.array(pred_t)
+
+                    if pred_t.size != 0:
+                        pred_class = np.array(pred_class)
+                        pred_class = pred_class[pred_t]
+
+                        pred_boxes = np.array(pred_boxes)
+                        pred_boxes = pred_boxes[pred_t]
+
+                        pred_score = np.array(pred_score)
+                        pred_score = pred_score[pred_t]
+
+                        outcome = outcome.append(
+                            {
+                                "labels": list(pred_class),
+                                "scores": list(pred_score),
+                                "boxes": list(pred_boxes)
+                            },
+                            ignore_index=True)
+
+
+        return outcome

src/udfs/pytorch_abstract_udf.py

@@ -38,8 +38,8 @@ def __init__(self):
         AbstractClassifierUDF.__init__(self)
         nn.Module.__init__(self)
 
-    def get_device(self):
-        return next(self.parameters()).device
+    #def get_device(self):
+        #return next(self.parameters()).device
 
     @property
     def transforms(self) -> Compose:
@@ -52,7 +52,7 @@ def transform(self, images: np.ndarray):
 
     def forward(self, frames: List[np.ndarray]):
         tens_batch = torch.cat([self.transform(x) for x in frames])\
-            .to(self.get_device())
+            .to(f"cuda:{self.device}")
         return self.classify(tens_batch)
 
     @abstractmethod

test/expression/test_function_expression.py

@@ -62,7 +62,7 @@ def test_function_move_the_device_to_gpu_if_compatible(self, context):
         gpu_mock_function = Mock(return_value=pd.DataFrame())
         gpu_device_id = '2'
 
-        mock_function.to_device.return_value = gpu_mock_function
+        mock_function.assign_device.return_value = gpu_mock_function
         context_instance.gpu_device.return_value = gpu_device_id
 
         expression = FunctionExpression(mock_function,
@@ -71,7 +71,7 @@ def test_function_move_the_device_to_gpu_if_compatible(self, context):
 
         input_batch = Batch(frames=pd.DataFrame())
         expression.evaluate(input_batch)
-        mock_function.to_device.assert_called_with(gpu_device_id)
+        mock_function.assign_device.assert_called_with(gpu_device_id)
         gpu_mock_function.assert_called()
 
     def test_should_use_the_same_function_if_not_gpu_compatible(self):

test/integration_tests/test_pytorch.py

@@ -69,3 +69,22 @@ def test_should_run_pytorch_and_ssd(self):
         res = actual_batch.frames
         for idx in res.index:
             self.assertTrue('car' in res['label'][idx])
+
+
+    def test_optical_character_recognition(self):
+        query = """LOAD DATA INFILE 'mnist.mp4'
+                   INTO MyVideo;"""
+        execute_query_fetch_all(query)
+
+        create_udf_query = """CREATE UDF IF NOT EXISTS OpticalCharacterRecognition
+                  INPUT  (Frame_Array NDARRAY UINT8(3, 256, 256))
+                  OUTPUT (label NDARRAY STR(10))
+                  TYPE  Classification
+                  IMPL  'src/udfs/optical_character_recognition.py';
+        """
+        execute_query_fetch_all(create_udf_query)
+
+        select_query = """SELECT OpticalCharacterRecognition(data) FROM MyVideo
+                        WHERE id <100;"""
+        actual_batch = execute_query_fetch_all(select_query)
+        self.assertEqual(actual_batch.batch_size, 42)