| 15 | logging.info('current time is {}'.format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))) |
| 16 | |
| 17 | class main(object): |
| 18 | def __init__(self, cfg, args): |
| 19 | self.classes, self.classes_num = load_classes(args.cls_file) |
| 20 | self.seg_model = load_seg_module(cfg) |
| 21 | self.test_data = load_test_data(cfg) |
| 22 | self.obj_model = OAM_GRAM(in_dim=1024, one_hot_cls_num=150).cuda().eval() |
| 23 | self.classifier = Classifier(num_classes=self.classes_num, in_dim=2048).cuda().eval() |
| 24 | if args.ckpt: |
| 25 | self.obj_model, self.classifier = \ |
| 26 | load_checkpoint(args.ckpt, self.obj_model, self.classifier) |
| 27 | self.cfg = cfg |
| 28 | self.args = args |
| 29 | self.correct = 0 |
| 30 | self.count = 0 |
| 31 | logging.info(self.obj_model) |
| 32 | logging.info(self.classifier) |
| 33 | |
| 34 | def get_object_feature(self, segSize, img_resized_list, batch_data): |
| 35 | segmentation_module = self.seg_model |
| 36 | args = self.args |
| 37 | cfg = self.cfg |
| 38 | |
| 39 | # Upload object features from local file instead of calculating the online |
| 40 | if args.local_object: |
| 41 | object_feature = get_obj_onehot_vector(batch_data['info'], args.local_object) |
| 42 | object_feature = torch.FloatTensor([object_feature]).cuda().view(1,1024,150,1) |
| 43 | return object_feature |
| 44 | |
| 45 | with torch.no_grad(): |
| 46 | scores = torch.zeros(1, cfg.DATASET.num_class, segSize[0], segSize[1]).cuda() |
| 47 | feature = torch.zeros(1, 1024, segSize[0], segSize[1]).cuda() |
| 48 | channels = feature.shape[1] |
| 49 | |
| 50 | for img in img_resized_list: |
| 51 | feed_dict = batch_data.copy() |
| 52 | feed_dict['img_data'] = img.cuda() |
| 53 | del feed_dict['img_ori'] |
| 54 | del feed_dict['info'] |
| 55 | pred_tmp, pred_tmp_feature_map = segmentation_module(feed_dict, segSize=segSize) |
| 56 | feature = feature + pred_tmp_feature_map / len(cfg.DATASET.imgSizes) |
| 57 | scores = scores + pred_tmp / len(cfg.DATASET.imgSizes) |
| 58 | |
| 59 | # Uncomment the following codes to verify the correctness of following codes. |
| 60 | # scores = torch.Tensor([[[[1, 0], [1, 0]], [[0, 1], [0, 0]], [[0, 0], [0, 1]]]]) |
| 61 | # feature = torch.Tensor([[[[0.5, 1], [0.5, 2]], [[0.5, 1], [0.5, 2]], [[0.5, 1], [0.5, 2]], [[0.5, 1], [0.5, 2]], [[0.5, 1], [0.5, 2]]]]) |
| 62 | # channels = 5 |
| 63 | # cfg.DATASET.num_class = 3 |
| 64 | |
| 65 | |
| 66 | # ***Open trigger will double the whole inference speed.**** |
| 67 | # ***However, this trigger will also slightly influence the object feature value** |
| 68 | # This inconsistency issue is stemmed from ***Pytorch***, not our method. |
| 69 | # Even the object feature value will slightly changed, our model is still stable. |
| 70 | trigger = True |
| 71 | # Object Feature Aggregation |
| 72 | if trigger: |
| 73 | scores = scores.view(cfg.DATASET.num_class, -1) |
| 74 | s, pred = torch.max(scores, dim=0) |