Machine Learning 之简单绘制决策树

#encoding:utf-8
#
#基本绘图
#
import numpy as np
import cv2
import cv2.cv as cv
class Parent:
      def __init__(self,k,d):
            self.k=k
            self.d=d

class DrawTree:
      def __init__(self,data):
            self.data=data
            self.canvas= np.zeros((600,1000,3),dtype="uint8")
            self.color={'green':(0,255,0),'red':(0,0,255),\
                        'blue':(255,0,0),'white':(255,255,255)}
            self.font=cv.InitFont(cv.CV_FONT_HERSHEY_PLAIN, 1.3, 1.3, 0, 1, 1)
            p1=Parent(0,50)
            self.getTree(p1,data,0,0)


      def getTree(self,parent,data,k,i):
            k=k+1
            for i in range(len(list(data))):
                  if type(data[list(data)[i]])==type({0:1}):
                        self.putText(str(data.keys()[i]),k*120,parent.d+i*300/k,'white')
                        parent=Parent(k*120,parent.d+i*350/k)
                        self.getTree(parent,data[list(data)[i]],k,i)   
                  else:
                      self.putText(str(list(data)[i]),(k)*120,parent.d+i*300/k,'white')
                      self.putText(str(data[list(data)[i]]),(k+1)*120,parent.d+i*300/k,'green')

      def putText(self,text,x,y,index):
            self.canvas=cv.fromarray(self.canvas)
            cv.PutText(self.canvas, text, (x-len(str(text))*10/2,y),self.font,self.color[index])
            self.canvas=np.array(self.canvas)
            #cv2.rectangle(self.canvas,(x-15,y-30),(x+30,y+15),self.color['green'])

      def putLine(self,x1,y1,x2,y2):
            cv2.line(self.canvas,(x1,x2),(y1,y2),self.color['red'],2)

      def showSave(self,name):
            cv2.imwrite(name,self.canvas)
            cv2.imshow("Canvas",self.canvas)
            cv2.waitKey(0)

测试 test.py

data={'no surfacing':{0:{'ash':{0:{0:'nnnnnn'},1:'yyyyyy'}},1:{'flippers':{0:'no',1:'yes'}}}}
a=DrawTree(data)
a.showSave('1.jpg')