Added legend and simplified do_classification