Live image/video translation
This post is part of the series On-device translation for Android
Using local translation models on Android
Adding OCR support to the local translator
Improving OCR quality on the translator
Live image/video translation (this article)
In the previous post, we got some high quality image translation, running reasonably fast, on-device.
We can push it further though, and avoid the need to take a picture, select it, and translate in discrete steps.
This post contains a bunch of jittery video, you might get motion-sick 🙃.
As a starting point, running a screenshot full of text on my phone takes:
Detection at 1400x3200: 500ms
Recognition: 1.5s
As usual, there are some easy tweaks that can give a large speed boosts.
First, the detection model scales linearly with pixel count and is quite resilient to noise. So we can shrink the image as much as possible1 and get it to run in ~80ms.
Then, the recognition model can be quantized. The published version uses fp32 for its calculations, and it can be quantized to either fp16 or int8, both of which have instructions that accelerate common operations at the hardware level (sdot and fp16). By quantizing to int8, the same screenshot executes in ~1200ms.
While this is a huge win for low effort, this brings us down to ~1.3s, which is not fast enough.
But that ~1.3s is for a screenshot packed with text. Recognition scales with the amount of text (obviously, more text = more to recognize), so for small sign it should finish in 50~100ms, which is good enough to try out a live overlay:
I spent a long while here, trying to match contours across frames but the heatmap from the detection model is not stable... and even if it was, inference is not fast enough to run on every frame.
So, we need a different way to track scene movement, which has the upside that we only need to run detection/recognition/translation once, on an initial "anchor" frame.
Given that we want to translate text, and text is usually on some kind of plane (a sign, paper, etc), we saw in the previous post that there's a single Homography that can express the exact transformation between two planes.
We can take this homography and apply it to the translated overlay, warping it the same way the scene has moved relative to the anchor.
How do we get there though?
Detect some image features using FAST, which is a "corner" detecting algorithm.
Then, describe the surroundings of each feature using rotated BRIEF, which is a way of converting the neighbors of the feature into a bunch of bits in a way that is rotation invariant and cheap to compare.
This gives us features + surroundings on the anchor, but some time has passed, and we now have a new, slightly different frame.
To match features between the two frames, we compare each of the new features to each of the original features, if we find a good match then we know how the single feature has transformed, otherwise, ignore the new feature.
From the surviving matches, we can apply RANSAC to fit a Homography: try matches at random and see how many agree to the proposed homography, keep the one with the most agreements.
The first version of this was almost like a miracle
It's still jittery though. Due to motion and sensor noise, FAST finds slightly different corners each frame, so each frame's homography gets fit from a different set of matches, and the resulting transform wobbles.
But we don't need to start from scratch on every frame, these frames are consecutive samples of a video stream, so we can use the temporal/spatial continuity between them to refine the feature search.
The Lucas-Kanade algorithm can carry confirmed matches from frame N to frame N+1 through optical flow ("seems like everything moved left, search for the features to the left").
Similarly, instead of computing a new homography per frame, we can use an Extended Kalman Filter to merge the new proposed H with the previous one, using different merge coefficients per degree of freedom (translation, perspective, rotation, ...).
With these new algorithms, the homography is stable over consecutive frames:
This works amazingly well, if a bit slow.
To make it faster, we have two phases to optimize, the single-shot detection/recognition/translation (inference), and the continuous homography tracking+overlay rendering.
Improving inference performance
We already discussed quantization and scaling, another easy win was to replace the translation engine (marian-nmt → slimt) which is ~4x faster, leaving recognition as the main bottleneck.
To improve the wall-clock time of the recognition model, we use 4 parallel sessions which work in batches; the problem here is that some lines are much wider than others, causing the whole batch to be delayed by the longest strip.
A good way to ensure no single strip is 'too wide' is to cap the max-width, by splitting on spaces. Spaces are cheap to compute: per dewarped strip, find lines perpendicular to reading...