diff --git a/GEODE/Visualisation/ConfusionMatrix.py b/GEODE/Visualisation/ConfusionMatrix.py index af140f30b697f8544382db92c4a43cc225b4114f..042ddcac941473a2d3255614841cde4dfce3c5a3 100644 --- a/GEODE/Visualisation/ConfusionMatrix.py +++ b/GEODE/Visualisation/ConfusionMatrix.py @@ -1,20 +1,11 @@ import argparse from GEODE.Store import prepare, tabular +from GEODE.Visualisation.Legend import trim import matplotlib.pyplot as plot import pandas import seaborn from sklearn.metrics import confusion_matrix -def trim(name, maxSize): - if len(name) > maxSize: - components = name.split(' ') - return components[0] + ' […]' - else: - return name - -def trimLabels(labels, maxWidth): - return labels if maxWidth is None else [trim(l, maxWidth) for l in labels] - def heatmap(matrix, filePath, labels, **kwargs): plot.figure(figsize=(16,13)) ax = seaborn.heatmap( @@ -46,7 +37,7 @@ def prepareData(data, labels=None): def drawConfusionMatrix(data, outputFile, labels=None, maxWidth=None, **kwargs): truth, answers, labels = prepareData(data, labels=labels) matrix = confusion_matrix(truth, answers, labels=labels, normalize='true') - heatmap(matrix, outputFile, trimLabels(labels, maxWidth), **kwargs) + heatmap(matrix, outputFile, trim(labels, maxWidth), **kwargs) def getArgs(arguments): cli = argparse.ArgumentParser( diff --git a/GEODE/Visualisation/__init__.py b/GEODE/Visualisation/__init__.py index 981a91757c5b5d4f6811f59844c12101ecc2d4f3..2eb7bce74fb7baa74ea7d1a36a095c0cfc315d43 100644 --- a/GEODE/Visualisation/__init__.py +++ b/GEODE/Visualisation/__init__.py @@ -1,2 +1,3 @@ from GEODE.Visualisation.ConfusionMatrix import drawConfusionMatrix, heatmap from GEODE.Visualisation.DensityProfile import densityProfile, drawDensityProfile, plotDensity +from GEODE.Visualisation.Legend import trim as legend diff --git a/GEODE/__init__.py b/GEODE/__init__.py index 563bce75422c75fc354eee084877a9af0ead7836..e7735f5fba302172f401b79b3b2fb6223885b756 100644 --- a/GEODE/__init__.py +++ b/GEODE/__init__.py @@ -23,7 +23,7 @@ from GEODE.ENE import eneLabels from GEODE.Metadata import article, articleKey, paragraph, paragraphKey, \ fromKey, relativePath, toKey, uid from GEODE.Store import corpus, Directory, SelfContained, tabular, toTSV -from GEODE.Visualisation import densityProfile, heatmap +from GEODE.Visualisation import densityProfile, heatmap, legend from GEODE.Visualisation.ConfusionMatrix import drawConfusionMatrixCLI from GEODE.Visualisation.DensityProfile import drawDensityProfileCLI