<reponame>dict/DeepMalwareLabeling<filename>backend/multitag/report.py
import pickle
import json
import numpy as np
import tensorflow as tf
import os
import csv
import random 

from sklearn.decomposition import PCA
import pandas as pd

from .estimator import MultilabelAddCenterlossClassifier
from .loader import ThumbnailMultilabelLoader
from .config import Config


def main():

    tf.logging.set_verbosity(tf.logging.INFO)
    config = Config()
    loader = ThumbnailMultilabelLoader()
    estimator = MultilabelAddCenterlossClassifier(loader.label_num).get_estimator()

    make_report(estimator, loader, config)


def make_report(estimator, loader, config):
    sha_train = get_sha_of_trainset(config.train_metadata_path)
    save_item(sha_train, config.sha_train_path)

    train_bottlenecks = get_bottlenecks_of_trainset(estimator, loader, config)
    save_item(train_bottlenecks, config.bottlenecks_train_path)

    valid_bottlenecks = get_bottlenecks_of_validset(estimator, loader, config)
    save_item(valid_bottlenecks, config.bottlenecks_valid_path)

    transformed, pca_object = get_pca(train_bottlenecks, 3)
    save_item(pca_object, config.pca_object_path)

    after_pca = get_after_pca(sha_train, transformed)
    save_item(after_pca, config.after_pca_path, 'json')

    tag_arr, tags_arr = get_tag_arr(config.train_metadata_path)
    pca3d_samples = get_pca3d_samples(tag_arr, tags_arr, transformed)
    save_item(pca3d_samples, config.pca3d_samples_path, 'json')

    centers = estimator.get_variable_value('centers')
    save_item(centers, config.centers_path)

    index_dict = get_index_dict(config.indexdict_path)
    save_item(config.tag_list_path, 'json')

def get_sha_of_trainset(train_metadata_path):
    sha_arr = []
    with open(train_metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            sha_arr.append(row['sha256'])
    return sha_arr


def get_index_dict(path):
    with open(path, 'rb') as f:
        index_dict = pickle.load(f)
    return index_dict

def get_bottlenecks_of_trainset(estimator, loader, config):
    eval_results = get_report_iterator(estimator, loader, config.train_metadata_path)
    return get_bottleneck_arr(eval_results)


def get_bottlenecks_of_validset(estimator, loader, config):
    eval_results = get_report_iterator(estimator, loader, config.valid_metadata_path)
    return get_bottleneck_arr(eval_results)


def get_bottleneck_arr(eval_results):
    bottleneck_arr = []
    for i, result in enumerate(eval_results):
        if i % 1000 == 0:
            print('%dth data is reported.'%i)
        bottleneck_arr.append(result['bottleneck'].tolist())

    return bottleneck_arr


def save_item(item, path, ftype='pickle'):
    if ftype == 'pickle':
        with open(path, 'wb') as f:
            pickle.dump(item, f)
    elif ftype == 'json':
        with open(path, 'w') as f:
            json.dump(item, f)
    else:
        raise NotImplementedError()


def get_report_iterator(estimator, loader, metadata_path):
    path_arr = []
    with open(metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            path = row['path']
            path_arr.append(path)

    eval_result = estimator.predict(input_fn=loader.predict_input_fn(path_arr))
    return eval_result


def get_pca(x, n_components=3):
    pca = PCA(n_components=n_components)
    transformed = pd.DataFrame(pca.fit_transform(x))
    return transformed, pca


def get_after_pca(sha_arr, transformed):
    after_pca = dict()
    transformed = np.array(transformed)
    for sha, point in zip(sha_arr, list(transformed)):
        after_pca[sha] = {'x':point[0], 'y':point[1], 'z':point[2]}

    return after_pca


def sort_tags(tags):
    splited_tags = tags.split(' ')
    sorted_tags = sorted(splited_tags)
    space = ' '
    str_tag = space.join(sorted_tags)
    return str_tag


def get_tag_arr(metadata_path, repr_idx=0):
    tag_arr = []
    tags_arr = []
    with open(metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            tag = random.choice(row['multitags'].split(' '))
            tags = sort_tags(row['multitags'])
            tag_arr.append(tag)
            tags_arr.append(tags)
    return np.array(tag_arr), np.array(tags_arr)


def get_pca3d_samples(tag_arr, tags_arr, transformed):
    keys = []
    for items in sorted(zip(*np.unique(tag_arr, return_counts=True)), key=lambda x: x[1],
                        reverse=True):
        keys.append(items[0])

    pca3d_samples = dict()
    text_samples = dict()
    for k in keys:
        pca3d_samples[k] = np.array(transformed[tag_arr==k]).tolist()
        text_samples[k] = list(map(lambda x: x[1], filter(lambda x: x[0], zip(tag_arr==k,tags_arr))))

    count = 0
    traces = []
    for k in pca3d_samples.keys():
        xl = []
        yl = []
        zl = []
        textl = text_samples[k]
        for i in pca3d_samples[k]:
            xl.append(i[0])
            yl.append(i[1])
            zl.append(i[2])
            count += 1
        trace = {
            'x':xl,
            'y':yl,
            'z':zl,
            'mode': 'markers',
            'marker': {
                'color': 'rgb({}, {}, {})'.format(random.randint(1, 255), random.randint(1, 255), random.randint(1, 255)),
                'size': 1,
            },
            'type': 'scatter3d',
            'name': k,
            'text': textl
        }
        traces.append(trace)
        if count >= 10000:
            break

    return traces
    



if __name__ == '__main__':
    main()


import pickle
import json
import numpy as np
import tensorflow as tf
import os
import csv
import random 

from sklearn.decomposition import PCA
import pandas as pd

from .estimator import MultilabelAddCenterlossClassifier
from .loader import ThumbnailMultilabelLoader
from .config import Config


def main():

    tf.logging.set_verbosity(tf.logging.INFO)
    config = Config()
    loader = ThumbnailMultilabelLoader()
    estimator = MultilabelAddCenterlossClassifier(loader.label_num).get_estimator()

    make_report(estimator, loader, config)


def make_report(estimator, loader, config):
    sha_train = get_sha_of_trainset(config.train_metadata_path)
    save_item(sha_train, config.sha_train_path)

    train_bottlenecks = get_bottlenecks_of_trainset(estimator, loader, config)
    save_item(train_bottlenecks, config.bottlenecks_train_path)

    valid_bottlenecks = get_bottlenecks_of_validset(estimator, loader, config)
    save_item(valid_bottlenecks, config.bottlenecks_valid_path)

    transformed, pca_object = get_pca(train_bottlenecks, 3)
    save_item(pca_object, config.pca_object_path)

    after_pca = get_after_pca(sha_train, transformed)
    save_item(after_pca, config.after_pca_path, 'json')

    tag_arr, tags_arr = get_tag_arr(config.train_metadata_path)
    pca3d_samples = get_pca3d_samples(tag_arr, tags_arr, transformed)
    save_item(pca3d_samples, config.pca3d_samples_path, 'json')

    centers = estimator.get_variable_value('centers')
    save_item(centers, config.centers_path)

    index_dict = get_index_dict(config.indexdict_path)
    save_item(config.tag_list_path, 'json')

def get_sha_of_trainset(train_metadata_path):
    sha_arr = []
    with open(train_metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            sha_arr.append(row['sha256'])
    return sha_arr


def get_index_dict(path):
    with open(path, 'rb') as f:
        index_dict = pickle.load(f)
    return index_dict

def get_bottlenecks_of_trainset(estimator, loader, config):
    eval_results = get_report_iterator(estimator, loader, config.train_metadata_path)
    return get_bottleneck_arr(eval_results)


def get_bottlenecks_of_validset(estimator, loader, config):
    eval_results = get_report_iterator(estimator, loader, config.valid_metadata_path)
    return get_bottleneck_arr(eval_results)


def get_bottleneck_arr(eval_results):
    bottleneck_arr = []
    for i, result in enumerate(eval_results):
        if i % 1000 == 0:
            print('%dth data is reported.'%i)
        bottleneck_arr.append(result['bottleneck'].tolist())

    return bottleneck_arr


def save_item(item, path, ftype='pickle'):
    if ftype == 'pickle':
        with open(path, 'wb') as f:
            pickle.dump(item, f)
    elif ftype == 'json':
        with open(path, 'w') as f:
            json.dump(item, f)
    else:
        raise NotImplementedError()


def get_report_iterator(estimator, loader, metadata_path):
    path_arr = []
    with open(metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            path = row['path']
            path_arr.append(path)

    eval_result = estimator.predict(input_fn=loader.predict_input_fn(path_arr))
    return eval_result


def get_pca(x, n_components=3):
    pca = PCA(n_components=n_components)
    transformed = pd.DataFrame(pca.fit_transform(x))
    return transformed, pca


def get_after_pca(sha_arr, transformed):
    after_pca = dict()
    transformed = np.array(transformed)
    for sha, point in zip(sha_arr, list(transformed)):
        after_pca[sha] = {'x':point[0], 'y':point[1], 'z':point[2]}

    return after_pca


def sort_tags(tags):
    splited_tags = tags.split(' ')
    sorted_tags = sorted(splited_tags)
    space = ' '
    str_tag = space.join(sorted_tags)
    return str_tag


def get_tag_arr(metadata_path, repr_idx=0):
    tag_arr = []
    tags_arr = []
    with open(metadata_path, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            tag = random.choice(row['multitags'].split(' '))
            tags = sort_tags(row['multitags'])
            tag_arr.append(tag)
            tags_arr.append(tags)
    return np.array(tag_arr), np.array(tags_arr)


def get_pca3d_samples(tag_arr, tags_arr, transformed):
    keys = []
    for items in sorted(zip(*np.unique(tag_arr, return_counts=True)), key=lambda x: x[1],
                        reverse=True):
        keys.append(items[0])

    pca3d_samples = dict()
    text_samples = dict()
    for k in keys:
        pca3d_samples[k] = np.array(transformed[tag_arr==k]).tolist()
        text_samples[k] = list(map(lambda x: x[1], filter(lambda x: x[0], zip(tag_arr==k,tags_arr))))

    count = 0
    traces = []
    for k in pca3d_samples.keys():
        xl = []
        yl = []
        zl = []
        textl = text_samples[k]
        for i in pca3d_samples[k]:
            xl.append(i[0])
            yl.append(i[1])
            zl.append(i[2])
            count += 1
        trace = {
            'x':xl,
            'y':yl,
            'z':zl,
            'mode': 'markers',
            'marker': {
                'color': 'rgb({}, {}, {})'.format(random.randint(1, 255), random.randint(1, 255), random.randint(1, 255)),
                'size': 1,
            },
            'type': 'scatter3d',
            'name': k,
            'text': textl
        }
        traces.append(trace)
        if count >= 10000:
            break

    return traces
    



if __name__ == '__main__':
    main()