def get_reward(filepath):
model_name = "model3D_1"
# create model
model = MultiColumn(config["num_classes"], cnn_def.Model,
int(config["column_units"]))
# multi GPU setting
model = torch.nn.DataParallel(model, device_ids).to(device)
save_dir = os.path.join(
"/scr1/system/beta-robot/smth-smth-v2-baseline-with-models/trained_models/pretrained/"
+ config['model_name'])
checkpoint_path = os.path.join(save_dir, 'model_best.pth.tar')
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
logits_matrix = []
features_matrix = []
targets_list = []
item_id_list = []
with torch.no_grad():
input = videoloader(filepath)
input = input.float().unsqueeze(0)
input_var = [input.to(device)]
output = model(input_var, False)
output = F.softmax(output, 1)
output = output.cpu().detach().numpy()
output = np.squeeze(output)
output_index = np.argsort(output * -1.0)
return (output, output_index)
nclips=config['nclips_val'],
step_size=config['step_size_val'],
is_val=True,
transform_pre=transform_eval_pre,
transform_post=transform_post,
get_item_id=True,
)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'],
pin_memory=True,
drop_last=False)
model.eval()
with torch.no_grad():
for i, (input, target, item_id) in enumerate(val_loader):
if config['nclips_val'] > 1:
input_var = list(input.split(config['clip_size'], 2))
for idx, inp in enumerate(input_var):
input_var[idx] = inp.to(device)
else:
input_var = [input.to(device)]
# compute output and loss
output, features = model(input_var, config['save_features'])
state = model.module.conv_column.h[0]
model_name = config["model_name"]
print("=> Name of the model -- {}".format(model_name))
# checkpoint path to a trained model
checkpoint_path = os.path.join("../", config["output_dir"], config["model_name"], "model_best.pth.tar")
print("=> Checkpoint path --> {}".format(checkpoint_path))
# ### 3. Load model
# In[7]:
model = MultiColumn(config['num_classes'], column_cnn_def.Model, int(config["column_units"]))
model.eval();
# In[8]:
print("=> loading checkpoint")
checkpoint = torch.load(checkpoint_path)
checkpoint['state_dict'] = remove_module_from_checkpoint_state_dict(
checkpoint['state_dict'])
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(checkpoint_path, checkpoint['epoch']))
# ### 4. Load data
class Base_eval:
def __init__(self, opt=None):
self.opt = opt
# self.root = '/scr1/system/beta-robot/base_eval'
# self.root = os.path.join(os.getcwd(),'base_eval')
if __name__ == '__main__':
self.root = '/scr1/system/gamma-robot/scripts/Eval/base_eval'
else:
self.root = os.path.join(self.opt.project_root, 'scripts', 'Eval',
'base_eval')
self.config = load_json_config(
os.path.join(self.root,
"configs/pretrained/config_model1_left_right.json"))
# set column model
# file_name = self.config['conv_model']
# self.cnn_def = importlib.import_module ("{}".format (file_name))
# setup device - CPU or GPU
self.device = torch.device("cuda")
self.device_ids = [0]
# model_name = "model3D_1"
# create model
# model = MultiColumn (self.config["num_classes"], self.cnn_def.Model, int (self.config["column_units"]))
self.model = MultiColumn(self.config["num_classes"], Model,
int(self.config["column_units"]))
# multi GPU setting
self.model = torch.nn.DataParallel(self.model,
self.device_ids).to(self.device)
try:
if self.opt.use_refine_baseline:
save_dir = os.path.join(
os.path.join(self.root,
"trained_models/pretrained/" + "refined"))
else:
save_dir = os.path.join(
os.path.join(
self.root, "trained_models/pretrained/" +
self.config['model_name']))
except:
save_dir = os.path.join(
os.path.join(
self.root,
"trained_models/pretrained/" + self.config['model_name']))
checkpoint_path = os.path.join(save_dir, 'model_best.pth.tar')
checkpoint = torch.load(checkpoint_path)
self.model.load_state_dict(checkpoint['state_dict'])
self.model.eval()
def videoloader(self, filepath):
transform_pre = ComposeMix([
[Scale(int(1.4 * self.config['input_spatial_size'])), "img"],
[torchvision.transforms.ToPILImage(), "img"],
[
torchvision.transforms.CenterCrop(
self.config['input_spatial_size']), "img"
],
])
transform_post = ComposeMix([
[torchvision.transforms.ToTensor(), "img"],
[
torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], # default values for imagenet
std=[0.229, 0.224, 0.225]),
"img"
]
])
imgs = []
for file in os.listdir(filepath):
file_id = int(file.split('.')[0])
tmp = cv2.imread(os.path.join(filepath, file))
tmp = cv2.cvtColor(tmp, cv2.COLOR_BGR2RGB)
imgs.append([tmp, file_id])
imgs = sorted(imgs, key=lambda x: x[1])
imgs = [x[0] for x in imgs]
imgs = transform_pre(imgs)
imgs = transform_post(imgs)
# num_frames = len (imgs)
# num_frames_necessary = 72
if len(imgs) < 72:
imgs.extend([imgs[-1]] * (72 - len(imgs)))
data = torch.stack(imgs)
data = data.permute(1, 0, 2, 3)
return data
def memroy_loader(self, imgs):
transform_pre = ComposeMix([
[Scale(int(1.4 * self.config['input_spatial_size'])), "img"],
[torchvision.transforms.ToPILImage(), "img"],
[
torchvision.transforms.CenterCrop(
self.config['input_spatial_size']), "img"
],
])
transform_post = ComposeMix([
[torchvision.transforms.ToTensor(), "img"],
[
torchvision.transforms.Normalize(
mean=[0.485, 0.456, 0.406], # default values for imagenet
std=[0.229, 0.224, 0.225]),
"img"
]
])
imgs = transform_pre(imgs)
imgs = transform_post(imgs)
# num_frames = len (imgs)
# num_frames_necessary = 72
if len(imgs) < 72:
imgs.extend([imgs[-1]] * (72 - len(imgs)))
data = torch.stack(imgs)
data = data.permute(1, 0, 2, 3)
return data
def get_baseline_reward(self, filepath):
with torch.no_grad():
input = self.videoloader(filepath)
input = input.float().unsqueeze(0)
input_var = [input.to(self.device)]
output = self.model(input_var, False)
output = F.softmax(output, 1)
output = output.cpu().detach().numpy()
output = np.squeeze(output)
output_index = np.argsort(output * -1.0)
return (output, output_index)
def get_memory_reward(self, img_buffer):
with torch.no_grad():
img_buffer = self.memroy_loader(img_buffer)
input = img_buffer.float().unsqueeze(0)
input_var = [input.to(self.device)]
output = self.model(input_var, False)
output = F.softmax(output, 1)
output = output.cpu().detach().numpy()
output = np.squeeze(output)
output_index = np.argsort(output * -1.0)
return (output, output_index)