Ocr code #197
Ocr code #197
Conversation
| outcome = pd.DataFrame() | ||
|
|
||
| for i in range(frames.shape[0]): | ||
| prediction = self.model.readtext_batched(frames[i]) |
There was a problem hiding this comment.
Isn't this logic processing one frame at a time? How about we call self.model.readtext_batched(frames)?
There was a problem hiding this comment.
I think there was some confusion due to the image size. I have made a few changes to make prediction on a batch of frames. Once we have the results for the entire batch, we traverse through each of the results and append them to the output dataframe. Creating a new PR for this!
| @@ -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): | |||
There was a problem hiding this comment.
Are these changes required? I remember we decided to revert these back. Please double-check, thanks!
There was a problem hiding this comment.
Yes, I think these changes are required. We are storing the device object using the assign_device() function in the GPU Compatible class as one of its variables. So we can directly access that object instead of using a get_device() function. I believe that was the reason we decided to eliminate this functionality. What do you think?
| @@ -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}") | |||
|
Resolved in #303 |
The code implementation for optical Character Recognition UDF.
A unit test has also been added in test_pytorch.py to test the OCR UDF.