Improve

just improve the codes.