Skip to content
Snippets Groups Projects
Commit f5e608cd authored by Mélodie Boillet's avatar Mélodie Boillet Committed by Yoann Schneider
Browse files

Add start_token parameter to prediction function

parent d8684909
No related branches found
No related tags found
1 merge request!94Add start_token parameter to prediction function
...@@ -92,6 +92,7 @@ class DAN: ...@@ -92,6 +92,7 @@ class DAN:
input_sizes, input_sizes,
confidences=False, confidences=False,
attentions=False, attentions=False,
start_token=None,
): ):
""" """
Run prediction on an input image. Run prediction on an input image.
...@@ -102,7 +103,9 @@ class DAN: ...@@ -102,7 +103,9 @@ class DAN:
""" """
input_tensor = input_tensor.to(self.device) input_tensor = input_tensor.to(self.device)
start_token = len(self.charset) + 1 start_token = (
self.charset.index(start_token) if start_token else len(self.charset) + 1
)
end_token = len(self.charset) end_token = len(self.charset)
# Run the prediction. # Run the prediction.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment