This repository introduces a reinforcement learning based weakly supervised system for localisation.
We train a controller function to localise regions of interest within an image by introducing a novel reward definition that utilises non-binarised classification probability, generated by a pre-trained binary classifier which classifies object presence in images or image crops. The object-presence classifier may then inform the controller of its localisation quality by quantifying the likelihood of the image containing an object. Such an approach allows us to minimize any potential labelling or human bias propagated via human labelling for fully supervised localisation.
The repository presents the code to train such a system (in train.py), with a saved model for prostate cancer localisation presented in predictor.
Examples of the predicted localisation, compared with other commonly used methods are presented below (green: ground truth; red: predicted):
