The current implementation can be improved by creating all classes beforehand and users can select whatever class to create training samples. https://user-images.githubusercontent.com/5016453/117547827-b994f380-afff-11eb-8371-144c4a778c54.mp4