import torch

import sys
sys.path.append('/workspace/home/jgusak/maxvol_objects/maskrcnn-benchmark/')
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.modeling.detector import build_detection_model, glue_detection_model

initial_dir = '/workspace/raid/data/jgusak/maxvol_objects/pretrained'
initial_path = '{}/facebook_frcnn_resnet50.pth'.format(initial_dir)

compressed_dir = '/workspace/raid/data/eponomarev/results/musco_voc/resnet50_fasterrcnn/tucker2/rank_selection:nx/ranks:/2x_layer_groups:1'

suffix = '_iter:{}-{}'.format(0, 1)
backbone_path = '{}/backbone{}.pth'.format(compressed_dir, suffix)
state_path = '{}/state{}.pth'.format(compressed_dir, suffix)

model = build_detection_model(cfg)

initial_backbone, rpn, head  = model.backbone, model.rpn, model.roi_heads

compressed_backbone = torch.load(backbone_path)
compressed_frcnn_model = glue_detection_model(compressed_backbone, rpn, head)

compressed_frcnn_model.eval()