import sys

sys.path.insert(0, '.')

from quiver_engine import server
from multiprocessing import Process
from quiver_engine.model_utils import register_hook
from searched_models.model_generator import model_builder


def visualize_feature_map(model_name):
    model, input_size = model_builder(model_name)

    hook_list = register_hook(model)

    server.launch(model, hook_list, input_folder="./data/Cat", image_size=input_size, use_gpu=False)


if __name__ == "__main__":
    p = Process(target=visualize_feature_map, args=('darts', ))
    p.start()
    while True:
        a = input("input:")
        p.terminate()
        p.join()
        print(a)
        if a == '0':
            p = Process(target=visualize_feature_map, args=('darts',))
            p.start()
        elif a == '1':
            p = Process(target=visualize_feature_map, args=('resnet18',))
            p.start()
        elif a == '2':
            p.terminate()
            p.join()
            break

