updated python script.

I made a couple of modifications on python script..