Python pydot 模块，graph_from_dot_data() 实例源码

def create_graph(filename, result_extract, verbose, path_write_png='/tmp/analysis_result.png', dangerous_score=5):
    #create DOT
    dot_content = 'digraph Analysis {\nratio=auto;\nnodesep="2.5 equally";\nranksep="2.5 equally";\n'
    color="green"
    if result_extract[u'GlobalRiskScore'] >= dangerous_score:
        color="red"
    dot_content += 'R_0 [shape=record, label="{{' + os.path.basename(filename) + '|' + str(result_extract[u'GlobalRiskScore']) + '|' + 'Coef:' + str(result_extract[u'GlobalRiskScoreCoef']) + '}|' + result_extract[u'RootFileType'].encode('utf8') + '}", color=' + color + '];\n'
    if result_extract[u'Yara']:
            dot_content += 'R_0_info [label="' + str(result_extract[u'Yara']).replace('}, {', '},\n{').replace('"', '').replace("'", '').encode('utf8') + '", color=blue];\n' 
            dot_content += 'R_0 -- R_0_info [style=dotted];\n'
    dot_content += json2dot(result_extract, dangerous_score, 'R_0', 'R_0')
    dot_content += '}'
    if verbose:
        print dot_content
    #convert dot to png
    (graph,) = pydot.graph_from_dot_data(dot_content)
    graph.write_png(path_write_png)

def decision_tree(X, y, regression, max_depth=3):
    from sklearn.tree import export_graphviz
    from sklearn.externals.six import StringIO  
    from IPython.core.pylabtools import figsize
    from IPython.display import Image
    figsize(12.5, 6)
    import pydot

    if regression:
        clf = DecisionTreeRegressor(max_depth=max_depth)
    else:
        clf = DecisionTreeClassifier(max_depth=max_depth)

    clf.fit(X, y)
    dot_data = StringIO()  
    export_graphviz(clf, out_file=dot_data, feature_names=list(X.columns),
                    filled=True, rounded=True,)
    graph = pydot.graph_from_dot_data(dot_data.getvalue())  
    return Image(graph.create_png())

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def read_hypergraph(string):
    """
    Read a hypergraph from a string in dot format. Nodes and edges specified in the input will be
    added to the current hypergraph.

    @type  string: string
    @param string: Input string in dot format specifying a graph.

    @rtype:  hypergraph
    @return: Hypergraph
    """
    hgr = hypergraph()
    dotG = pydot.graph_from_dot_data(string)

    # Read the hypernode nodes...
    # Note 1: We need to assume that all of the nodes are listed since we need to know if they
    #           are a hyperedge or a normal node
    # Note 2: We should read in all of the nodes before putting in the links
    for each_node in dotG.get_nodes():
        if 'hypernode' == each_node.get('hyper_node_type'):
            hgr.add_node(each_node.get_name())
        elif 'hyperedge' == each_node.get('hyper_node_type'):
            hgr.add_hyperedge(each_node.get_name())

    # Now read in the links to connect the hyperedges
    for each_link in dotG.get_edges():
        if hgr.has_node(each_link.get_source()):
            link_hypernode = each_link.get_source()
            link_hyperedge = each_link.get_destination()
        elif hgr.has_node(each_link.get_destination()):
            link_hypernode = each_link.get_destination()
            link_hyperedge = each_link.get_source()
        hgr.link(link_hypernode, link_hyperedge)

    return hgr

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def parse_elf(workspace, file):
    r2 = r2pipe.open(file.filepath)
    r2.cmd("aa")
    r2.cmd("afl")
    result = r2.cmd("agC")
    output_dir = os.path.join(workspace, "graphs")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    out_file = os.path.join(output_dir, file.hash)
    graph = pydot.graph_from_dot_data(result)
    graph[0].write_png(out_file)
    file.graph_file = out_file
    file.save()
    print("%s parsed" % file.filepath)

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def method2format(output, _format="png", mx=None, raw=None):
    """
        Export method to a specific file format

        @param output : output filename
        @param _format : format type (png, jpg ...) (default : png)
        @param mx : specify the MethodAnalysis object
        @param raw : use directly a dot raw buffer if None
    """
    try:
        import pydot
    except ImportError:
        error("module pydot not found")

    buff = "digraph {\n"
    buff += "graph [rankdir=TB]\n"
    buff += "node [shape=plaintext]\n"

    if raw:
        data = raw
    else:
        data = method2dot(mx)

    # subgraphs cluster
    buff += "subgraph cluster_" + hashlib.md5(output).hexdigest() + " {\nlabel=\"%s\"\n" % data['name']
    buff += data['nodes']
    buff += "}\n"

    # subgraphs edges
    buff += data['edges']
    buff += "}\n"

    d = pydot.graph_from_dot_data(buff)
    if d:
        getattr(d, "write_" + _format.lower())(output)

def classifyTree(Xtr, ytr, Xte, yte, splitCriterion="gini", maxDepth=0, visualizeTree=False):
    """ Classifies data using CART """
    try:
        accuracyRate, probabilities, timing = 0.0, [], 0.0
        # Perform classification
        cartClassifier = tree.DecisionTreeClassifier(criterion=splitCriterion, max_depth=maxDepth)
        startTime = time.time()
        prettyPrint("Training a CART tree for classification using \"%s\" and maximum depth of %s" % (splitCriterion, maxDepth), "debug")
        cartClassifier.fit(numpy.array(Xtr), numpy.array(ytr))
        prettyPrint("Submitting the test samples", "debug")
        predicted = cartClassifier.predict(Xte)
        endTime = time.time()
        # Compare the predicted and ground truth and append result to list
        accuracyRate = round(metrics.accuracy_score(predicted, yte), 2)
        # Also append the probability estimates
        probs = cartClassifier.predict_proba(Xte)
        probabilities.append(probs)
        timing = endTime-startTime # Keep track of performance
        if visualizeTree:
            # Visualize the tree
            dot_data = StringIO()
            tree.export_graphviz(cartClassifier, out_file=dot_data)
            graph = pydot.graph_from_dot_data(dot_data.getvalue())
            prettyPrint("Saving learned CART to \"tritonTree_%s.pdf\"" % getTimestamp(), "debug")
            graph.write_pdf("tree_%s.pdf" % getTimestamp())

    except Exception as e:
        prettyPrint("Error encountered in \"classifyTree\": %s" % e, "error")

    return accuracyRate, timing, probabilities, predicted

def visualize_tree(clf, outname, headers):
    from sklearn.externals.six import StringIO
    import pydot
    dot_data = StringIO()
    tree.export_graphviz(clf, out_file=dot_data, feature_names=list(headers))
    graph = pydot.graph_from_dot_data(dot_data.getvalue().decode('latin1').encode('utf8'))
    graph.write_pdf(outname)

def read(string):
    """
    Read a graph from a string in Dot language and return it. Nodes and edges specified in the
    input will be added to the current graph.

    @type  string: string
    @param string: Input string in Dot format specifying a graph.

    @rtype: graph
    @return: Graph
    """

    dotG = pydot.graph_from_dot_data(string)

    if (dotG.get_type() == "graph"):
        G = graph()
    elif (dotG.get_type() == "digraph"):
        G = digraph()
    elif (dotG.get_type() == "hypergraph"):
        return read_hypergraph(string)
    else:
        raise InvalidGraphType

    # Read nodes...
    # Note: If the nodes aren't explicitly listed, they need to be
    for each_node in dotG.get_nodes():
        G.add_node(each_node.get_name())
        for each_attr_key, each_attr_val in each_node.get_attributes().items():
            G.add_node_attribute(each_node.get_name(), (each_attr_key, each_attr_val))

    # Read edges...
    for each_edge in dotG.get_edges():
        # Check if the nodes have been added
        if not G.has_node(each_edge.get_source()):
            G.add_node(each_edge.get_source())
        if not G.has_node(each_edge.get_destination()):
            G.add_node(each_edge.get_destination())

        # See if there's a weight
        if 'weight' in each_edge.get_attributes().keys():
            _wt = each_edge.get_attributes()['weight']
        else:
            _wt = 1

        # See if there is a label
        if 'label' in each_edge.get_attributes().keys():
            _label = each_edge.get_attributes()['label']
        else:
            _label = ''

        G.add_edge((each_edge.get_source(), each_edge.get_destination()), wt = _wt, label = _label)

        for each_attr_key, each_attr_val in each_edge.get_attributes().items():
            if not each_attr_key in ['weight', 'label']:
                G.add_edge_attribute((each_edge.get_source(), each_edge.get_destination()), \
                                            (each_attr_key, each_attr_val))

    return G

def makePrediction(para,rawData,totalNumRows,labels):
    traingSetSize=int(math.floor(totalNumRows*para['trainingSetPercent']))
    print('%d instances are selected as training dataset!'%traingSetSize)
    trainX=np.array(rawData[0:traingSetSize])
    trainY=np.array(labels[0:traingSetSize])
    clf=tree.DecisionTreeClassifier()
    clf=clf.fit(trainX,trainY)

    feaNames=['event'+str(i) for i in range(1,386)]
    classNames=trainY
    # generate the decision tree figure
    # dot_data = StringIO()  #class_names=classNames,
    # tree.export_graphviz(clf, out_file=dot_data, feature_names=feaNames,
    #                      filled=True, rounded=True,
    #                      special_characters=True)
    # graph = pydot.graph_from_dot_data(dot_data.getvalue())
    # graph.write_png('sample_SOSP.png')

    testingX=rawData[traingSetSize:]
    testingY=labels[traingSetSize:]
    prediction=list(clf.predict(testingX))
    if len(prediction)!=len(testingY):
        print ('prediction and testingY have different length and SOMEWHERE WRONG!')
    sameLabelNum=0
    sameFailureNum=0
    for i in range(len(testingY)):
        if prediction[i]==testingY[i]:
            sameLabelNum+=1
            if prediction[i]==1:
                sameFailureNum+=1

    accuracy=float(sameLabelNum)/len(testingY)
    print ('accuracy is %.5f:'%accuracy)

    predictSuccess=0
    predictFailure=0
    for item in prediction:
        if item==0:
            predictSuccess+=1
        elif item==1:
            predictFailure+=1

    testSuccess=0
    testFailure=0
    for tt in testingY:
        if tt==0:
            testSuccess+=1
        elif tt==1:
            testFailure+=1

    print(predictSuccess,predictFailure,testSuccess,testFailure,sameFailureNum)
    if sameFailureNum==0:
        print ('precision is 0 and recall is 0')
    else:
        precision=float(sameFailureNum)/(predictFailure)
        print('precision is %.5f'%precision)
        recall=float(sameFailureNum)/(testFailure)
        print('recall is %.5f'%recall)
        F_measure=2*precision*recall/(precision+recall)
        print('F_measure is %.5f'%F_measure)
    return predictFailure,testFailure,sameFailureNum,precision,recall,F_measure

def classifyTreeKFold(X, y, kFold=2, splitCriterion="gini", maxDepth=0, visualizeTree=False):
    """ Classifies data using CART and K-Fold cross validation """
    try:
        groundTruthLabels, predictedLabels = [], []
        accuracyRates = [] # Meant to hold the accuracy rates
        # Split data into training and test datasets
        trainingDataset, testDataset = [], []
        trainingLabels, testLabels = [], []
        accuracyRates = []
        probabilities = []
        timings = []
        kFoldValidator = KFold(n=len(X), n_folds=kFold, shuffle=False)
        currentFold = 1
        for trainingIndices, testIndices in kFoldValidator:
            # Prepare the training and testing datasets
            for trIndex in trainingIndices:
                trainingDataset.append(X[trIndex])
                trainingLabels.append(y[trIndex])
            for teIndex in testIndices:
                testDataset.append(X[teIndex])
                testLabels.append(y[teIndex])
            # Perform classification
            startTime = time.time()
            cartClassifier = tree.DecisionTreeClassifier(criterion=splitCriterion, max_depth=maxDepth)
            prettyPrint("Training a CART tree for classification using \"%s\" and maximum depth of %s" % (splitCriterion, maxDepth), "debug")
            cartClassifier.fit(numpy.array(trainingDataset), numpy.array(trainingLabels))
            prettyPrint("Submitting the test samples", "debug")
            predicted = cartClassifier.predict(testDataset)
            endTime = time.time()
            # Add that to the groundTruthLabels and predictedLabels matrices
            groundTruthLabels.append(testLabels)
            predictedLabels.append(predicted)
            # Compare the predicted and ground truth and append result to list
            accuracyRates.append(round(metrics.accuracy_score(predicted, testLabels), 2))
            # Also append the probability estimates
            probs = cartClassifier.predict_proba(testDataset)
            probabilities.append(probs)
            timings.append(endTime-startTime) # Keep track of performance
            if visualizeTree:
                # Visualize the tree
                dot_data = StringIO() 
                tree.export_graphviz(cartClassifier, out_file=dot_data) 
                graph = pydot.graph_from_dot_data(dot_data.getvalue()) 
                prettyPrint("Saving learned CART to \"tritonTree_%s.pdf\"" % currentFold, "debug")
                graph.write_pdf("tritonTree_%s.pdf" % currentFold)

            trainingDataset, trainingLabels = [], []
            testDataset, testLabels = [], []
            currentFold += 1 

    except Exception as e:
        prettyPrint("Error encountered in \"classifyTreeKFold\": %s" % e, "error")
        return [], [], []

    return accuracyRates, probabilities, timings, groundTruthLabels, predictedLabels