update

Author: we1pingyu

cal4od/cal4od_helper.py

@@ -268,7 +268,7 @@ def intersect(boxes1, boxes2):
 import matplotlib.pyplot as plt
 
 
-def draw_PIL_image(image, boxes, labels, scores, name, no=None, color='greenyellow'):
+def draw_PIL_image(image, boxes, labels, scores, name):
     if type(image) != PIL.Image.Image:
         image = F.to_pil_image(image)
     plt.imshow(image)
@@ -279,20 +279,15 @@ def draw_PIL_image(image, boxes, labels, scores, name, no=None, color='greenyell
     plt.margins(0, 0)
     plt.gca().xaxis.set_major_locator(plt.NullLocator())
     plt.gca().yaxis.set_major_locator(plt.NullLocator())
-    i = 0
-    if no is not None:
-        for n in no:
-            if i >= 1:
-                color = 'greenyellow'
-            else:
-                color = 'red'
-            i += 1
-            x, y = boxes[n][0], boxes[n][1]
-            w, h = boxes[n][2] - boxes[n][0], boxes[n][3] - boxes[n][1]
-            plt.gca().add_patch(plt.Rectangle((x, y), w, h, fill=False, edgecolor=color, linewidth=2.5))
-            plt.text(x, y, '{}={}'.format(voc_labels[labels[n]], scores[n]), color='color', verticalalignment='bottom',
-                     fontsize=4)
-    plt.savefig('vis/{}.jpg'.format(name), dpi=128, bbox_inches='tight', pad_inches=0)
+    # for i in range(len(boxes)):
+    #     x, y = boxes[i][0], boxes[i][1]
+    #     w, h = boxes[i][2] - boxes[i][0], boxes[i][3] - boxes[i][1]
+    #     plt.gca().add_patch(
+    #         plt.Rectangle((x, y), w, h, fill=False, edgecolor=label_color_map[rev_label_map[labels[i].item()]],
+    #                       linewidth=2.5))
+    #     # plt.text(x, y, '{}={}'.format(voc_labels[labels[n]], scores[n]), color='color', verticalalignment='bottom',
+    #     #              fontsize=4)
+    plt.savefig('vis/{}.png'.format(name), dpi=256, bbox_inches='tight', pad_inches=0)
     # plt.show()
     plt.cla()
 

cal4od_train.py

@@ -36,6 +36,8 @@
 from detection.frcnn_la import fasterrcnn_resnet50_fpn_feature
 from detection.retinanet_cal import retinanet_mobilenet, retinanet_resnet50_fpn_cal
 
+A = 1
+
 
 def train_one_epoch(task_model, task_optimizer, data_loader, device, cycle, epoch, print_freq):
     task_model.train()
@@ -53,6 +55,10 @@ def train_one_epoch(task_model, task_optimizer, data_loader, device, cycle, epoc
     for images, targets in metric_logger.log_every(data_loader, print_freq, header):
         images = list(image.to(device) for image in images)
         targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        # global A
+        # for image, target in zip(images, targets):
+        #     draw_PIL_image(image, target['boxes'], target['labels'], None, "unlabeled_{}".format(A))
+        #     A += 1
         task_loss_dict = task_model(images, targets)
         task_losses = sum(loss for loss in task_loss_dict.values())
         # reduce losses over all GPUs for logging purposes
@@ -203,8 +209,8 @@ def get_uncertainty(task_model, unlabeled_loader, augs, num_cls):
                 for output, aug_box, aug_image in zip(outputs, aug_boxes, aug_images):
                     consistency_img = 1.0
                     mean_img = []
-                    boxes, scores_cls, pm, labels = output['boxes'], output['scores_cls'], output['prob_max'], output[
-                        'labels']
+                    boxes, scores_cls, pm, labels, scores = output['boxes'], output['scores_cls'], output['prob_max'], \
+                                                            output['labels'], output['scores']
                     cls_corr = [0] * (num_cls - 1)
                     for p, l in zip(pm, labels):
                         # if p.item() > 0.4:
@@ -217,7 +223,7 @@ def get_uncertainty(task_model, unlabeled_loader, augs, num_cls):
                         continue
                     j = 0
                     no = []
-                    for ab, ref_score_cls, ref_pm in zip(aug_box, ref_scores_cls, prob_max):
+                    for ab, ref_score_cls, ref_pm, ref_score in zip(aug_box, ref_scores_cls, prob_max, ref_scores):
                         width = torch.min(ab[2], boxes[:, 2]) - torch.max(ab[0], boxes[:, 0])
                         height = torch.min(ab[3], boxes[:, 3]) - torch.max(ab[1], boxes[:, 1])
                         Aarea = (ab[2] - ab[0]) * (ab[3] - ab[1])
@@ -232,13 +238,13 @@ def get_uncertainty(task_model, unlabeled_loader, augs, num_cls):
                         js = 0.5 * scipy.stats.entropy(p, m) + 0.5 * scipy.stats.entropy(q, m)
                         if js < 0:
                             js = 0
-                        # if 0.6 > torch.max(iou) > 0.4 and 0.4 < 0.5 * (1 - js) * (
-                        #         ref_pm + pm[torch.argmax(iou)]) < 0.6 and (ab[2] - ab[0]) > 200 and \
-                        #         torch.max(ref_score_cls) < 0.6 and pm[torch.argmax(iou)] < 0.6:
+                        # if 0.7 < torch.max(iou) and 0.4 > 0.5 * (1 - js) * (
+                        #         ref_pm + pm[torch.argmax(iou)]) and (ab[2] - ab[0]) > 100 and scores[
+                        #     torch.argmax(iou)] > 0.4 and ref_score > 0.4:
                         #     # draw_PIL_image(image, boxes, ref_labels, i, no=[ab], color='greenyellow')
                         #     no = [j, torch.argmax(iou).item()]
                         #     draw_PIL_image_2(aug_image.cpu(), aug_box, boxes, ref_labels, labels, ref_scores, pm,
-                        #                      'bad', no=no, color='red')
+                        #                      'bad_2', no=no, color='red')
                         #     print(1 / 0)
                         j += 1
                         consistency_img = min(consistency_img, torch.abs(
@@ -379,7 +385,7 @@ def main(args):
         init_num = 500
         budget_num = 500
         if 'retina' in args.model:
-            init_num = 500
+            init_num = 1000
             budget_num = 500
     else:
         init_num = 5000
@@ -441,14 +447,23 @@ def main(args):
             print("Getting stability")
             random.shuffle(unlabeled_set)
             if 'coco' in args.dataset:
-                subset = unlabeled_set[:10000]
+                subset = unlabeled_set[:5000]
             else:
                 subset = unlabeled_set
             if args.mutual:
                 unlabeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(subset),
                                               num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn)
                 uncertainty, _cls_corrs = get_uncertainty(task_model, unlabeled_loader, augs, num_classes)
+                # labeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(labeled_set),
+                #                             num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn)
+                # u, _ = get_uncertainty(task_model, labeled_loader, augs, num_classes)
+                # with open("vis/cal_labeled_metric_{}_{}_{}.pkl".format(args.model, args.dataset, cycle),
+                #           "wb") as fp:  # Pickling
+                #     pickle.dump(u, fp)
                 arg = np.argsort(np.array(uncertainty))
+                # with open("vis/cal_unlabeled_metric_{}_{}_{}.pkl".format(args.model, args.dataset, cycle),
+                #           "wb") as fp:  # Pickling
+                #     pickle.dump(torch.tensor(uncertainty)[arg][:budget_num].numpy(), fp)
                 cls_corrs_set = arg[:int(args.mr * budget_num)]
                 cls_corrs = [_cls_corrs[i] for i in cls_corrs_set]
                 labeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(labeled_set),
@@ -457,8 +472,6 @@ def main(args):
                 # Update the labeled dataset and the unlabeled dataset, respectively
                 tobe_labeled_set = list(torch.tensor(subset)[arg][tobe_labeled_set].numpy())
                 labeled_set += tobe_labeled_set
-                with open("vis/cal_{}_{}_{}.txt".format(args.model, args.dataset, cycle), "wb") as fp:  # Pickling
-                    pickle.dump(labeled_set, fp)
                 unlabeled_set = list(set(indices) - set(labeled_set))
             else:
                 unlabeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(subset),
@@ -467,6 +480,7 @@ def main(args):
                 arg = np.argsort(np.array(uncertainty))
                 # Update the labeled dataset and the unlabeled dataset, respectively
                 labeled_set += list(torch.tensor(subset)[arg][:budget_num].numpy())
+                labeled_set = list(set(labeled_set))
                 unlabeled_set = list(set(indices) - set(labeled_set))
 
             # Create a new dataloader for the updated labeled dataset
@@ -494,26 +508,35 @@ def main(args):
                     coco_evaluate(task_model, data_loader_test)
                 elif 'voc' in args.dataset:
                     voc_evaluate(task_model, data_loader_test, args.dataset, False, path=args.results_path)
-        if not args.skip and cycle == 0:
-            if 'faster' in args.model:
-                utils.save_on_master({
-                    'model': task_model.state_dict(), 'args': args},
-                    os.path.join(args.first_checkpoint_path, '{}_frcnn_1st.pth'.format(args.dataset)))
-            elif 'retina' in args.model:
-                utils.save_on_master({
-                    'model': task_model.state_dict(), 'args': args},
-                    os.path.join(args.first_checkpoint_path, '{}_retinanet_1st.pth'.format(args.dataset)))
+        # if not args.skip and cycle == 0:
+        #     if 'faster' in args.model:
+        #         utils.save_on_master({
+        #             'model': task_model.state_dict(), 'args': args},
+        #             os.path.join(args.first_checkpoint_path, '{}_frcnn_1st.pth'.format(args.dataset)))
+        #     elif 'retina' in args.model:
+        #         utils.save_on_master({
+        #             'model': task_model.state_dict(), 'args': args},
+        #             os.path.join(args.first_checkpoint_path, '{}_retinanet_1st.pth'.format(args.dataset)))
         random.shuffle(unlabeled_set)
         if 'coco' in args.dataset:
-            subset = unlabeled_set[:10000]
+            subset = unlabeled_set[:5000]
         else:
             subset = unlabeled_set
         print("Getting stability")
         if args.mutual:
             unlabeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(subset),
                                           num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn)
             uncertainty, _cls_corrs = get_uncertainty(task_model, unlabeled_loader, augs, num_classes)
+            labeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(labeled_set),
+                                        num_workers=args.workers, pin_memory=True, collate_fn=utils.collate_fn)
+            # u, _ = get_uncertainty(task_model, labeled_loader, augs, num_classes)
+            # with open("vis/cal_labeled_metric_{}_{}_{}.pkl".format(args.model, args.dataset, cycle),
+            #           "wb") as fp:  # Pickling
+            #     pickle.dump(u, fp)
             arg = np.argsort(np.array(uncertainty))
+            # with open("vis/cal_unlabeled_metric_{}_{}_{}.pkl".format(args.model, args.dataset, cycle),
+            #           "wb") as fp:  # Pickling
+            #     pickle.dump(torch.tensor(uncertainty)[arg][:budget_num].numpy(), fp)
             cls_corrs_set = arg[:int(args.mr * budget_num)]
             cls_corrs = [_cls_corrs[i] for i in cls_corrs_set]
             labeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(labeled_set),
@@ -522,8 +545,8 @@ def main(args):
             # Update the labeled dataset and the unlabeled dataset, respectively
             tobe_labeled_set = list(torch.tensor(subset)[arg][tobe_labeled_set].numpy())
             labeled_set += tobe_labeled_set
-            with open("vis/cal_{}_{}_{}.txt".format(args.model, args.dataset, cycle), "wb") as fp:  # Pickling
-                pickle.dump(labeled_set, fp)
+            # with open("vis/cal_{}_{}_{}.txt".format(args.model, args.dataset, cycle), "wb") as fp:  # Pickling
+            #     pickle.dump(labeled_set, fp)
             unlabeled_set = list(set(indices) - set(labeled_set))
         else:
             unlabeled_loader = DataLoader(dataset_aug, batch_size=1, sampler=SubsetSequentialSampler(subset),
@@ -532,6 +555,7 @@ def main(args):
             arg = np.argsort(np.array(uncertainty))
             # Update the labeled dataset and the unlabeled dataset, respectively
             labeled_set += list(torch.tensor(subset)[arg][:budget_num].numpy())
+            labeled_set = list(set(labeled_set))
             unlabeled_set = list(set(indices) - set(labeled_set))
         # Create a new dataloader for the updated labeled dataset
         train_sampler = SubsetRandomSampler(labeled_set)

detection/retina_ssm.py

@@ -522,27 +522,34 @@ def ssm_postprocess_detections(self, head_outputs, anchors, image_shapes):
             if torch.max(scores_per_image) < CONF_THRESH:
                 # print(scores)
                 al_idx = 1
+                detections.append({
+                    "boxes": all_boxes,
+                    "labels": all_labels,
+                    "scores": all_scores,
+                    'al': al_idx,
+                })
                 continue
             for class_index in range(num_classes):
                 # remove low scoring boxes
                 inds = torch.gt(scores_per_image[:, class_index], self.score_thresh)
-                boxes_per_class, scores_per_class, labels_per_class = \
-                    boxes_per_image[inds], scores_per_image[inds, class_index], labels_per_image[inds, class_index]
+                boxes_per_class, scores_per_class, scores_all_class, labels_per_class = \
+                    boxes_per_image[inds], scores_per_image[inds, class_index], scores_per_image[inds], \
+                    labels_per_image[inds, class_index]
                 other_outputs_per_class = [(k, v[inds]) for k, v in other_outputs_per_image]
 
                 # remove empty boxes
                 keep = box_ops.remove_small_boxes(boxes_per_class, min_size=1e-2)
-                boxes_per_class, scores_per_class, labels_per_class = \
-                    boxes_per_class[keep], scores_per_class[keep], labels_per_class[keep]
+                boxes_per_class, scores_per_class, scores_all_class, labels_per_class = \
+                    boxes_per_class[keep], scores_per_class[keep], scores_all_class[keep], labels_per_class[keep]
                 other_outputs_per_class = [(k, v[keep]) for k, v in other_outputs_per_class]
 
                 # non-maximum suppression, independently done per class
                 keep = box_ops.nms(boxes_per_class, scores_per_class, self.nms_thresh)
 
                 # keep only topk scoring predictions
                 keep = keep[:self.detections_per_img]
-                boxes_per_class, scores_per_class, labels_per_class = \
-                    boxes_per_class[keep], scores_per_class[keep], labels_per_class[keep]
+                boxes_per_class, scores_per_class, scores_all_class, labels_per_class = \
+                    boxes_per_class[keep], scores_per_class[keep], scores_all_class[keep], labels_per_class[keep]
                 other_outputs_per_class = [(k, v[keep]) for k, v in other_outputs_per_class]
 
                 image_boxes.append(boxes_per_class)
@@ -553,19 +560,20 @@ def ssm_postprocess_detections(self, head_outputs, anchors, image_shapes):
                     if k not in image_other_outputs:
                         image_other_outputs[k] = []
                     image_other_outputs[k].append(v)
-                for i in len(boxes_per_class):
+
+                for i in range(len(boxes_per_class)):
                     all_boxes = torch.cat((all_boxes, boxes_per_class[i].unsqueeze(0)), 0)
                     all_scores = torch.cat((all_scores, scores_per_class[i].unsqueeze(0)), 0)
-                    all_labels.append(judge_y(scores_per_class[i]))
-
+                    all_labels.append(judge_y(scores_all_class[i][1:]))
+            detections.append({
+                "boxes": all_boxes,
+                "labels": all_labels,
+                "scores": all_scores,
+                'al': al_idx,
+            })
             for k, v in image_other_outputs.items():
                 detections[-1].update({k: torch.cat(v, dim=0)})
-        detections.append({
-            "boxes": all_boxes,
-            "labels": all_labels,
-            "scores": all_scores,
-            'al': al_idx,
-        })
+
         return detections
 
     def forward(self, images, targets=None):
@@ -646,7 +654,6 @@ def forward(self, images, targets=None):
             # print(self.ssm)
             if self.ssm:
                 detections = self.ssm_postprocess_detections(head_outputs, anchors, images.image_sizes)
-                print(detections)
             else:
                 detections = self.postprocess_detections(head_outputs, anchors, images.image_sizes)
             detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)