Minor fixes to python script
Added appropriate labels