update readme

WenmuZhou · WenmuZhou · commit 561c1fe31b9d · 2019-10-24T12:28:09.000+08:00
diff --git a/README.MD b/README.MD
@@ -60,7 +60,7 @@ only train on ICDAR2015 dataset
 | Method                   | image size (short size) |learning rate | Precision (%) | Recall (%) | F-measure (%) | FPS |
 |--------------------------|-------|--------|--------|------------|---------------|-----|
 | paper(resnet18)  | 736 |x | x | x | 80.4 | 26.1 |
-| my (resnet18+FPEM_FFM+pse扩张)  |736 |1e-3| 73.17 | 70.39 | 71.75 | 21.31 (P100)|
+| my (resnet18+FPEM_FFM+pse扩张)  |736 |1e-3| 84.24 | 74.14 | 78.87 | 21.31 (P100)|
 | my (resnet50+FPEM_FFM+pse扩张)  |736 |1e-3| 69.04 | 66.66 | 67.83 | 14.22 (P100)|
 | my (resnet18+FPEM_FFM+pse扩张)  |736 |1e-4| 62.93 | 62.41 | 62.61 | 21.31 (P100)|
 | my (resnet50+FPEM_FFM+pse扩张)  |736 |1e-4| 61.19 | 69.18 | 64.94 | 14.22 (P100)|
diff --git a/base/base_trainer.py b/base/base_trainer.py
@@ -20,7 +20,7 @@ def __init__(self, config, model, criterion, weights_init):
         self.save_dir = os.path.join(config['trainer']['output_dir'], config['name'])
         self.checkpoint_dir = os.path.join(self.save_dir, 'checkpoint')
 
-        if config['trainer']['resume']['restart_training']:
+        if config['trainer']['resume_checkpoint'] == '' and config['trainer']['finetune_checkpoint'] == '':
             shutil.rmtree(self.save_dir, ignore_errors=True)
         if not os.path.exists(self.checkpoint_dir):
             os.makedirs(self.checkpoint_dir)
@@ -42,7 +42,7 @@ def __init__(self, config, model, criterion, weights_init):
         self.logger = setup_logger(os.path.join(self.save_dir, 'train_log'))
         self.logger.info(pformat(self.config))
 
-        # device set
+        # device
         torch.manual_seed(self.config['trainer']['seed'])  # 为CPU设置随机种子
         if len(self.config['trainer']['gpus']) > 0 and torch.cuda.is_available():
             self.with_cuda = True
@@ -62,8 +62,10 @@ def __init__(self, config, model, criterion, weights_init):
 
         self.optimizer = self._initialize('optimizer', torch.optim, model.parameters())
 
-        if self.config['trainer']['resume']['checkpoint'] != '':
-            self._resume_checkpoint(self.config['trainer']['resume']['checkpoint'])
+        if self.config['trainer']['resume_checkpoint'] != '':
+            self._laod_checkpoint(self.config['trainer']['resume_checkpoint'], resume=True)
+        elif self.config['trainer']['finetune_checkpoint'] != '':
+            self._laod_checkpoint(self.config['trainer']['finetune_checkpoint'], resume=False)
         else:
             if weights_init is not None:
                 model.apply(weights_init)
@@ -171,15 +173,15 @@ def _save_checkpoint(self, epoch, file_name, save_best=False):
         else:
             self.logger.info("Saving checkpoint: {}".format(filename))
 
-    def _resume_checkpoint(self, resume_path):
+    def _laod_checkpoint(self, checkpoint_path, resume):
         """
         Resume from saved checkpoints
-        :param resume_path: Checkpoint path to be resumed
+        :param checkpoint_path: Checkpoint path to be resumed
         """
-        self.logger.info("Loading checkpoint: {} ...".format(resume_path))
-        checkpoint = torch.load(resume_path)
+        self.logger.info("Loading checkpoint: {} ...".format(checkpoint_path))
+        checkpoint = torch.load(checkpoint_path)
         self.model.load_state_dict(checkpoint['state_dict'])
-        if not self.config['trainer']['resume']['restart_training']:
+        if resume:
             self.global_step = checkpoint['global_step']
             self.start_epoch = checkpoint['epoch'] + 1
             self.config['lr_scheduler']['args']['last_epoch'] = self.start_epoch
@@ -192,8 +194,9 @@ def _resume_checkpoint(self, resume_path):
                     for k, v in state.items():
                         if isinstance(v, torch.Tensor):
                             state[k] = v.to(self.device)
-        # self.config = checkpoint['config']
-        self.logger.info("Checkpoint '{}' (epoch {}) loaded".format(resume_path, self.start_epoch))
+            self.logger.info("resume from checkpoint {} (epoch {})".format(checkpoint_path, self.start_epoch))
+        else:
+            self.logger.info("finetune from checkpoint {}".format(checkpoint_path))
 
     def _initialize(self, name, module, *args, **kwargs):
         module_name = self.config[name]['type']
diff --git a/config.json b/config.json
@@ -77,10 +77,8 @@
         "epochs": 600,
         "display_interval": 10,
         "show_images_interval": 50,
-        "resume": {
-            "restart_training": true,
-            "checkpoint": ""
-        },
+        "resume_checkpoint": "",
+        "finetune_checkpoint": "",
         "output_dir": "output",
         "tensorboard": true
     }
diff --git a/data_loader/dataset.py b/data_loader/dataset.py
@@ -6,7 +6,7 @@
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 from data_loader.data_utils import image_label
-from utils import order_points_colckwise
+from utils import order_points_clockwise
 
 
 class ImageDataset(Dataset):
@@ -52,7 +52,7 @@ def _get_annotation(self, label_path: str) -> tuple:
             for line in f.readlines():
                 params = line.strip().strip('\ufeff').strip('\xef\xbb\xbf').split(',')
                 try:
-                    box = order_points_colckwise(np.array(list(map(float, params[:8]))).reshape(-2, 1))
+                    box = order_points_clockwise(np.array(list(map(float, params[:8]))).reshape(-1, 2))
                     if cv2.arcLength(box, True) > 0:
                         boxes.append(box)
                         label = params[8]
diff --git a/eval.py b/eval.py
@@ -41,9 +41,9 @@ def main(model_path, img_folder, save_path, gpu_id):
 
 
 if __name__ == '__main__':
-    os.environ['CUDA_VISIBLE_DEVICES'] = str('2')
+    os.environ['CUDA_VISIBLE_DEVICES'] = str('0')
     scale = 4
-    model_path = 'output/PAN_pred_mask_resnet50/checkpoint/model_best.pth'
+    model_path = 'output/PAN_resnet18_FPEM_FFM/checkpoint/model_best.pth'
     img_path = '/data2/dataset/ICD15/test/img'
     gt_path = '/data2/dataset/ICD15/test/gt'
     save_path = model_path.replace('checkpoint/model_best.pt', 'result/')
diff --git a/models/model.py b/models/model.py
@@ -61,10 +61,10 @@ def forward(self, x):
         'fpem_repeat': 4,  # fpem模块重复的次数
         'pretrained': False,  # backbone 是否使用imagesnet的预训练模型
         'result_num':7,
-        'segmentation_head': 'FPEM_FFM'  # 分割头，FPN or FPEM
+        'segmentation_head': 'FPEM_FFM'  # 分割头，FPN or FPEM_FFM
     }
     model = Model(model_config=model_config).to(device)
     y = model(x)
     print(y.shape)
-    print(model)
+    # print(model)
     # torch.save(model.state_dict(), 'PAN.pth')
diff --git a/models/modules/segmentation_head.py b/models/modules/segmentation_head.py
@@ -16,42 +16,40 @@ def __init__(self, backbone_out_channels, **kwargs):
         result_num = kwargs.get('result_num', 6)
         inplace = True
         conv_out = 256
-
-        # Top layer
-        self.toplayer = nn.Sequential(
-            nn.Conv2d(backbone_out_channels[3], conv_out, kernel_size=1, stride=1, padding=0),
+        # reduce layers
+        self.reduce_conv_c2 = nn.Sequential(
+            nn.Conv2d(backbone_out_channels[0], conv_out, kernel_size=1, stride=1, padding=0),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-        # Lateral layers
-        self.latlayer1 = nn.Sequential(
-            nn.Conv2d(backbone_out_channels[2], conv_out, kernel_size=1, stride=1, padding=0),
+        self.reduce_conv_c3 = nn.Sequential(
+            nn.Conv2d(backbone_out_channels[1], conv_out, kernel_size=1, stride=1, padding=0),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-        self.latlayer2 = nn.Sequential(
-            nn.Conv2d(backbone_out_channels[1], conv_out, kernel_size=1, stride=1, padding=0),
+        self.reduce_conv_c4 = nn.Sequential(
+            nn.Conv2d(backbone_out_channels[2], conv_out, kernel_size=1, stride=1, padding=0),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-        self.latlayer3 = nn.Sequential(
-            nn.Conv2d(backbone_out_channels[0], conv_out, kernel_size=1, stride=1, padding=0),
+
+        self.reduce_conv_c5 = nn.Sequential(
+            nn.Conv2d(backbone_out_channels[3], conv_out, kernel_size=1, stride=1, padding=0),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-
         # Smooth layers
-        self.smooth1 = nn.Sequential(
+        self.smooth_p4 = nn.Sequential(
             nn.Conv2d(conv_out, conv_out, kernel_size=3, stride=1, padding=1),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-        self.smooth2 = nn.Sequential(
+        self.smooth_p3 = nn.Sequential(
             nn.Conv2d(conv_out, conv_out, kernel_size=3, stride=1, padding=1),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
         )
-        self.smooth3 = nn.Sequential(
+        self.smooth_p2 = nn.Sequential(
             nn.Conv2d(conv_out, conv_out, kernel_size=3, stride=1, padding=1),
             nn.BatchNorm2d(conv_out),
             nn.ReLU(inplace=inplace)
@@ -67,27 +65,27 @@ def __init__(self, backbone_out_channels, **kwargs):
     def forward(self, x):
         c2, c3, c4, c5 = x
         # Top-down
-        p5 = self.toplayer(c5)
-        p4 = self._upsample_add(p5, self.latlayer1(c4))
-        p4 = self.smooth1(p4)
-        p3 = self._upsample_add(p4, self.latlayer2(c3))
-        p3 = self.smooth2(p3)
-        p2 = self._upsample_add(p3, self.latlayer3(c2))
-        p2 = self.smooth3(p2)
+        p5 = self.reduce_conv_c5(c5)
+        p4 = self._upsample_add(p5, self.reduce_conv_c4(c4))
+        p4 = self.smooth_p4(p4)
+        p3 = self._upsample_add(p4, self.reduce_conv_c3(c3))
+        p3 = self.smooth_p3(p3)
+        p2 = self._upsample_add(p3, self.reduce_conv_c2(c2))
+        p2 = self.smooth_p2(p2)
 
         x = self._upsample_cat(p2, p3, p4, p5)
         x = self.conv(x)
         x = self.out_conv(x)
         return x
 
     def _upsample_add(self, x, y):
-        return F.interpolate(x, size=y.size()[2:], mode='bilinear', align_corners=False) + y
+        return F.interpolate(x, size=y.size()[2:], mode='bilinear') + y
 
     def _upsample_cat(self, p2, p3, p4, p5):
         h, w = p2.size()[2:]
-        p3 = F.interpolate(p3, size=(h, w), mode='bilinear', align_corners=False)
-        p4 = F.interpolate(p4, size=(h, w), mode='bilinear', align_corners=False)
-        p5 = F.interpolate(p5, size=(h, w), mode='bilinear', align_corners=False)
+        p3 = F.interpolate(p3, size=(h, w), mode='bilinear')
+        p4 = F.interpolate(p4, size=(h, w), mode='bilinear')
+        p5 = F.interpolate(p5, size=(h, w), mode='bilinear')
         return torch.cat([p2, p3, p4, p5], dim=1)
 
 
@@ -99,22 +97,40 @@ def __init__(self, backbone_out_channels, **kwargs):
         """
         super().__init__()
         fpem_repeat = kwargs.get('fpem_repeat', 2)
-        self.conv_c2 = nn.Conv2d(in_channels=backbone_out_channels[0], out_channels=128, kernel_size=1)
-        self.conv_c3 = nn.Conv2d(in_channels=backbone_out_channels[1], out_channels=128, kernel_size=1)
-        self.conv_c4 = nn.Conv2d(in_channels=backbone_out_channels[2], out_channels=128, kernel_size=1)
-        self.conv_c5 = nn.Conv2d(in_channels=backbone_out_channels[3], out_channels=128, kernel_size=1)
+        conv_out = 128
+        # reduce layers
+        self.reduce_conv_c2 = nn.Sequential(
+            nn.Conv2d(in_channels=backbone_out_channels[0], out_channels=conv_out, kernel_size=1),
+            nn.BatchNorm2d(conv_out),
+            nn.ReLU()
+        )
+        self.reduce_conv_c3 = nn.Sequential(
+            nn.Conv2d(in_channels=backbone_out_channels[1], out_channels=conv_out, kernel_size=1),
+            nn.BatchNorm2d(conv_out),
+            nn.ReLU()
+        )
+        self.reduce_conv_c4 = nn.Sequential(
+            nn.Conv2d(in_channels=backbone_out_channels[2], out_channels=conv_out, kernel_size=1),
+            nn.BatchNorm2d(conv_out),
+            nn.ReLU()
+        )
+        self.reduce_conv_c5 = nn.Sequential(
+            nn.Conv2d(in_channels=backbone_out_channels[3], out_channels=conv_out, kernel_size=1),
+            nn.BatchNorm2d(conv_out),
+            nn.ReLU()
+        )
         self.fpems = nn.ModuleList()
         for i in range(fpem_repeat):
-            self.fpems.append(FPEM(128))
-        self.out_conv = nn.Conv2d(in_channels=512, out_channels=6, kernel_size=1)
+            self.fpems.append(FPEM(conv_out))
+        self.out_conv = nn.Conv2d(in_channels=conv_out * 4, out_channels=6, kernel_size=1)
 
     def forward(self, x):
         c2, c3, c4, c5 = x
         # reduce channel
-        c2 = self.conv_c2(c2)
-        c3 = self.conv_c3(c3)
-        c4 = self.conv_c4(c4)
-        c5 = self.conv_c5(c5)
+        c2 = self.reduce_conv_c2(c2)
+        c3 = self.reduce_conv_c3(c3)
+        c4 = self.reduce_conv_c4(c4)
+        c5 = self.reduce_conv_c5(c5)
 
         # FPEM
         for i, fpem in enumerate(self.fpems):
@@ -142,38 +158,28 @@ def forward(self, x):
 class FPEM(nn.Module):
     def __init__(self, in_channels=128):
         super().__init__()
-        # self.add_up = nn.Sequential(
-        #     nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3, padding=1, groups=in_channel),
-        #     nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=1),
-        #     nn.BatchNorm2d(in_channel),
-        #     nn.ReLU()
-        # )
-        # self.add_down = nn.Sequential(
-        #     nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=3, padding=1, groups=in_channel,
-        #               stride=2),
-        #     nn.Conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=1),
-        #     nn.BatchNorm2d(in_channel),
-        #     nn.ReLU()
-        # )
         self.up_add1 = SeparableConv2d(in_channels, in_channels, 1)
         self.up_add2 = SeparableConv2d(in_channels, in_channels, 1)
-        self.up_add2 = SeparableConv2d(in_channels, in_channels, 1)
+        self.up_add3 = SeparableConv2d(in_channels, in_channels, 1)
         self.down_add1 = SeparableConv2d(in_channels, in_channels, 2)
         self.down_add2 = SeparableConv2d(in_channels, in_channels, 2)
         self.down_add3 = SeparableConv2d(in_channels, in_channels, 2)
 
     def forward(self, c2, c3, c4, c5):
         # up阶段
-        c4 = self.up_add1(c4 + F.interpolate(c5, c4.size()[-2:], mode='bilinear', align_corners=True))
-        c3 = self.up_add2(c3 + F.interpolate(c4, c3.size()[-2:], mode='bilinear', align_corners=True))
-        c2 = self.up_add2(c2 + F.interpolate(c3, c2.size()[-2:], mode='bilinear', align_corners=True))
+        c4 = self.up_add1(self._upsample_add(c5, c4))
+        c3 = self.up_add2(self._upsample_add(c4, c3))
+        c2 = self.up_add3(self._upsample_add(c3, c2))
 
         # down 阶段
-        c3 = self.down_add1(c2 + F.interpolate(c3, c2.size()[-2:], mode='bilinear', align_corners=True))
-        c4 = self.down_add2(c3 + F.interpolate(c4, c3.size()[-2:], mode='bilinear', align_corners=True))
-        c5 = self.down_add3(c4 + F.interpolate(c5, c4.size()[-2:], mode='bilinear', align_corners=True))
+        c3 = self.down_add1(self._upsample_add(c3, c2))
+        c4 = self.down_add2(self._upsample_add(c4, c3))
+        c5 = self.down_add3(self._upsample_add(c5, c4))
         return c2, c3, c4, c5
 
+    def _upsample_add(self, x, y):
+        return F.interpolate(x, size=y.size()[2:], mode='bilinear') + y
+
 
 class SeparableConv2d(nn.Module):
     def __init__(self, in_channels, out_channels, stride=1):
@@ -186,8 +192,8 @@ def __init__(self, in_channels, out_channels, stride=1):
         self.relu = nn.ReLU()
 
     def forward(self, x):
-        x = self.out_channels(x)
+        x = self.depthwise_conv(x)
         x = self.pointwise_conv(x)
         x = self.bn(x)
-        x = self.relu
+        x = self.relu(x)
         return x
diff --git a/post_processing/__init__.py b/post_processing/__init__.py
@@ -25,7 +25,7 @@ def decode(preds, scale=1, threshold=0.7311, min_area=5):
     :param threshold: sigmoid的阈值
     :return: 最后的输出图和文本框
     """
-    from .pse import pse_cpp, get_points, get_sum
+    from .pse import pse_cpp, get_points, get_num
     preds[:2, :, :] = torch.sigmoid(preds[:2, :, :])
     preds = preds.detach().cpu().numpy()
     score = preds[0].astype(np.float32)
@@ -35,42 +35,28 @@ def decode(preds, scale=1, threshold=0.7311, min_area=5):
 
     label_num, label = cv2.connectedComponents(kernel.astype(np.uint8), connectivity=4)
     label_values = []
-    label_sum = get_sum(label, label_num)
+    label_sum = get_num(label, label_num)
     for label_idx in range(1, label_num):
         if label_sum[label_idx] < min_area:
             continue
         label_values.append(label_idx)
 
     pred = pse_cpp(text.astype(np.uint8), similarity_vectors, label, label_num, 0.8)
     pred = pred.reshape(text.shape)
-    bbox_list = []
-    for label_value in label_values:
-        points = np.array(np.where(pred == label_value)).transpose((1, 0))[:, ::-1]
-
-        if points.shape[0] < 100 / (scale * scale):
-            continue
-
-        score_i = np.mean(score[pred == label_value])
-        if score_i < 0.1:
-            continue
-
-        rect = cv2.minAreaRect(points)
-        bbox = cv2.boxPoints(rect)
-        bbox_list.append([bbox[1], bbox[2], bbox[3], bbox[0]])
 
     bbox_list = []
     label_points = get_points(pred, score, label_num)
-    for label_value, label_point in label_points.item():
+    for label_value, label_point in label_points.items():
         if label_value not in label_values:
             continue
         score_i = label_point[0]
         label_point = label_point[2:]
-        points = np.array(label_point).reshape(-1, 2)
+        points = np.array(label_point, dtype=int).reshape(-1, 2)
 
         if points.shape[0] < 100 / (scale * scale):
             continue
 
-        if score_i < 0.1:
+        if score_i < 0.93:
             continue
 
         rect = cv2.minAreaRect(points)
diff --git a/utils/cal_recall/rrc_evaluation_funcs.py b/utils/cal_recall/rrc_evaluation_funcs.py
diff --git a/utils/util.py b/utils/util.py
