code cleanup

Changed some of the imports. Just importing the methods needed instead
of the whole package.