import torch
from torchviz import make_dot
import os


def visual_model(model, inputs, framework='hikdet', file_dir='./'):
    x = inputs
    if framework == 'pytorch':
        y = model(x)
    elif framework == 'hikdet':
        y = model.forward_dummy(x)
    else:
        raise NotImplementedError
    g = make_dot(y[0])
    file_path = os.path.join(file_dir, 'model_graph')
    g.render(filename=file_path, view=False, format='png')
    return file_path+'.png'
