romainbrette / holypipette

Holy Pipette
GNU General Public License v3.0
4 stars 0 forks source link

Yolo tracking #103

Open mstimberg opened 2 years ago

mstimberg commented 2 years ago

This adds support for tracking of Paramecia with a yolov5 network, based on the work by @maya-awada. There are various things that are a bit hacky, but in general it works (note that yolov5 needs Python 3, so Python2-only devices like the current MicroManager module cannot be used with it).

The bulk of the implementation is in the YoloTracker class in holypipette.interface.tracker. Since yolov5 is not packaged, we need to point to its directory and add it to sys.path. YoloTracker therefore requires this path as an argument – it should be stored in the yolo_path variable in setup_script.py. The GUI implementation is in TrackerGui in holypipette.gui.tracker, but it only adds a single button (👁) and the configuration to the standard Camera GUI. To start everything, run the yolo_gui.py script in the main repository.

The integration of the detection into holypipette is very simple, but probably not the most optimal. The detection method is handed over to the Camera GUI as an "image post-processing" (image_edit) method, and the drawing of the bounding boxes uses the "display overlaid information (display_edit) method (see https://holypipette.readthedocs.io/en/latest/developer.html#adding-functionality). This means that when you activate the detection, it will block the display of each camera image until it has done the detection/annotation. The camera display will therefore look "choppy". It might be nicer to instead do the detection in a separate parallel thread, and continue to display all camera images. But then the overlaid bounding boxes would probably no longer match the paramecia that were detected earlier.

The image is cropped to a given size (see configuration tab), and only the cropped image is handed over to yolov5. The cropped size has to match the size that the network was trained with, I did not handle any scaling. By default, it displays the full image with a box for the cropped area, but you can also choose to only display the cropped area. The detected Paramecium that is the closest to the center of the crop area is annotated with a bounding box in yellow, any additional Paramecia in the image get a blue bounding box.

As a first step towards moving the stage to follow the Paramecium, the configuration already has a "move stage" option, but all this actually does is to move around the crop area. If you only display the cropped area, this does look as if the stage is following, though :wink: But this could be directly replaced by a call to move the stage, in particular if it is a non-blocking call that sends off the command and returns immediately. The main reason why I implemented it with a move in the crop area instead is that this was the only thing I could do in my test video...