HRNet (heatmap+classification)

Author: nuguziii

experiments/pano/hrnet/pano_test1.yaml

@@ -37,7 +37,7 @@
 _C.MODEL.NAME = 'pose_hrnet'
 _C.MODEL.INIT_WEIGHTS = True
 _C.MODEL.PRETRAINED = ''
-_C.MODEL.NUM_JOINTS = 17
+_C.MODEL.NUM_JOINTS = 32
 _C.MODEL.TAG_PER_JOINT = True
 _C.MODEL.TARGET_TYPE = 'gaussian'
 _C.MODEL.IMAGE_SIZE = [256, 256]  # width * height, ex: 192 * 256
@@ -82,6 +82,7 @@
 _C.TRAIN.NESTEROV = False
 _C.TRAIN.GAMMA1 = 0.99
 _C.TRAIN.GAMMA2 = 0.0
+_C.TRAIN.LOSS_WEIGHT = [1e1,1e1,1e1]
 
 _C.TRAIN.BEGIN_EPOCH = 0
 _C.TRAIN.END_EPOCH = 140

lib/config/default.py

@@ -70,4 +70,45 @@ def accuracy(output, target, hm_type='gaussian', thr=0.5):
         acc[0] = avg_acc
     return acc, avg_acc, cnt, pred
 
+def accuracy_classification(output, target, thres=0.0):
+    '''
+    Calculate accuracy according to PCK,
+    but uses ground truth heatmap rather than x,y locations
+    First value to be returned is average accuracy across 'idxs',
+    followed by individual accuracies
+    '''
+
+    output[output<thres] = 0
+    output[output>=thres] = 1
+
+    cnt = output.shape[0]
+    acc = np.zeros(cnt)
+
+    target = target.reshape(cnt, -1)
+
+    for i in range(cnt):
+        acc[i] = sum(output[i]==target[i])/32
+
+    avg_acc = np.mean(acc)
+    return avg_acc, cnt
+
+def accuracy_landmark(output, target):
+    '''
+    Calculate accuracy according to PCK,
+    but uses ground truth heatmap rather than x,y locations
+    First value to be returned is average accuracy across 'idxs',
+    followed by individual accuracies
+    '''
+
+    cnt = output.shape[0]
+    acc = np.zeros(cnt)
+
+    target = target.reshape(cnt, -1)
+
+    for i in range(cnt):
+        acc[i] = sum(output[i]==target[i])/32
+
+    avg_acc = np.mean(acc)
+    return avg_acc, cnt
+
 

lib/core/evaluate.py

@@ -15,7 +15,7 @@
 import numpy as np
 import torch
 
-from core.evaluate import accuracy
+from core.evaluate import accuracy, accuracy_classification, accuracy_landmark
 from core.inference import get_final_preds
 from utils.transforms import flip_back
 from utils.vis import save_result_images, save_debug_images
@@ -29,6 +29,9 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
     batch_time = AverageMeter()
     data_time = AverageMeter()
     losses = AverageMeter()
+    loss_classifier = AverageMeter()
+    loss_heatmap = AverageMeter()
+    loss_landmark = AverageMeter()
     acc = AverageMeter()
 
     # switch to train mode
@@ -40,19 +43,27 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
         data_time.update(time.time() - end)
 
         # compute output
-        outputs = model(input)
+        heatmap, classification, landmark = model(input)
 
         target = target.cuda(non_blocking=True)
         target_weight = target_weight.cuda(non_blocking=True)
 
-        if isinstance(outputs, list):
-            loss = criterion(outputs[0], target, target_weight)
-            for output in outputs[1:]:
-                loss += criterion(output, target, target_weight)
+        if isinstance(heatmap, list):
+            heatloss = criterion[0](heatmap[0], target, target_weight)
+            for output in heatmap[1:]:
+                heatloss += criterion[0](output, target, target_weight)
         else:
-            output = outputs
-            loss = criterion(output, target, target_weight)
+            output = heatmap
+            heatloss = criterion[0](output, target, target_weight)
 
+        #target2 = meta["visible"].type(torch.FloatTensor).cuda(non_blocking=True).view(classification.size(0),-1)
+        target2 = meta["visible"].type(torch.FloatTensor).cuda(non_blocking=True)
+        classloss = criterion[1](classification, target2)
+
+        target3 = meta["joints"].reshape(-1,64).type(torch.FloatTensor).cuda(non_blocking=True)
+        lmloss = criterion[2](landmark, target3)
+
+        loss = config.TRAIN.LOSS_WEIGHT[1]*classloss + config.TRAIN.LOSS_WEIGHT[2]*lmloss
         # loss = criterion(output, target, target_weight)
 
         # compute gradient and do update step
@@ -62,6 +73,9 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
 
         # measure accuracy and record loss
         losses.update(loss.item(), input.size(0))
+        loss_classifier.update(classloss.item(), input.size(0))
+        loss_heatmap.update(heatloss.item(), input.size(0))
+        loss_landmark.update(lmloss.item(), input.size(0))
 
         _, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
                                          target.detach().cpu().numpy())
@@ -76,19 +90,15 @@ def train(config, train_loader, model, criterion, optimizer, epoch,
                   'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
                   'Speed {speed:.1f} samples/s\t' \
                   'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
-                  'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
-                  'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
+                  'Loss {loss.val:.5f} ({loss.avg:.5f}) ({classific.avg: .5f}+{lm.avg: .5f})\t' \
+                  'Accuracy(heatmap) {acc.val:.3f} ({acc.avg:.3f})'.format(
                       epoch, i, len(train_loader), batch_time=batch_time,
                       speed=input.size(0)/batch_time.val,
-                      data_time=data_time, loss=losses, acc=acc)
+                      data_time=data_time,
+                      loss=losses, classific=loss_classifier, lm=loss_landmark,
+                      acc=acc)
             logger.info(msg)
 
-            writer = writer_dict['writer']
-            global_steps = writer_dict['train_global_steps']
-            writer.add_scalar('train_loss', losses.val, global_steps)
-            writer.add_scalar('train_acc', acc.val, global_steps)
-            writer_dict['train_global_steps'] = global_steps + 1
-
             prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
             save_debug_images(config, input, meta, target, pred*4, output,
                               prefix)
@@ -117,11 +127,11 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
         end = time.time()
         for i, (input, target, target_weight, meta) in enumerate(val_loader):
             # compute output
-            outputs = model(input)
-            if isinstance(outputs, list):
-                output = outputs[-1]
+            heatmap, classification, landmark = model(input)
+            if isinstance(heatmap, list):
+                output = heatmap[-1]
             else:
-                output = outputs
+                output = heatmap
 
             if config.TEST.FLIP_TEST:
                 input_flipped = input.flip(3)
@@ -147,7 +157,11 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             target = target.cuda(non_blocking=True)
             target_weight = target_weight.cuda(non_blocking=True)
 
-            loss = criterion(output, target, target_weight)
+            target2 = meta["visible"].type(torch.FloatTensor).cuda(non_blocking=True)
+            target3 = meta["joints"].reshape(-1, 64).type(torch.FloatTensor).cuda(non_blocking=True)
+
+            loss = config.TRAIN.LOSS_WEIGHT[1]*criterion[1](classification, target2) \
+                 + config.TRAIN.LOSS_WEIGHT[2] * criterion[2](landmark, target3)
 
             num_images = input.size(0)
             # measure accuracy and record loss
@@ -201,6 +215,7 @@ def test(config, val_loader, val_dataset, model, criterion, output_dir,
     batch_time = AverageMeter()
     losses = AverageMeter()
     acc = AverageMeter()
+    acc_clas = AverageMeter()
 
     # switch to evaluate mode
     model.eval()
@@ -219,11 +234,11 @@ def test(config, val_loader, val_dataset, model, criterion, output_dir,
         end = time.time()
         for i, (input, target, target_weight, meta) in enumerate(val_loader):
             # compute output
-            outputs = model(input)
-            if isinstance(outputs, list):
-                output = outputs[-1]
+            heatmap, classification = model(input)
+            if isinstance(heatmap, list):
+                output = heatmap[-1]
             else:
-                output = outputs
+                output = heatmap
 
             if config.TEST.FLIP_TEST:
                 input_flipped = input.flip(3)
@@ -249,16 +264,21 @@ def test(config, val_loader, val_dataset, model, criterion, output_dir,
             target = target.cuda(non_blocking=True)
             target_weight = target_weight.cuda(non_blocking=True)
 
-            loss = criterion(output, target, target_weight)
+            target_class = meta["visible"].type(torch.FloatTensor).cuda(non_blocking=True)
+
+            loss = config.TRAIN.LOSS_WEIGHT[0]*criterion[0](output, target, target_weight) + criterion[1](classification,target_class)
 
             num_images = input.size(0)
             # measure accuracy and record loss
             losses.update(loss.item(), num_images)
             _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
                                              target.cpu().numpy())
-
             acc.update(avg_acc, cnt)
 
+            avg_acc, cnt = accuracy_classification(classification.cpu().numpy(),
+                                                   target_class.cpu().numpy())
+            acc_clas.update(avg_acc, cnt)
+
             # measure elapsed time
             batch_time.update(time.time() - end)
             end = time.time()
@@ -285,9 +305,10 @@ def test(config, val_loader, val_dataset, model, criterion, output_dir,
                 msg = 'Test: [{0}/{1}]\t' \
                       'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
-                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
+                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})\t'\
+                      'Accuracy {acc2.val:.3f} ({acc2.avg:.3f})'.format(
                           i, len(val_loader), batch_time=batch_time,
-                          loss=losses, acc=acc)
+                          loss=losses, acc=acc, acc2=acc_clas)
                 logger.info(msg)
 
                 prefix = os.path.join(output_dir, 'result')

lib/core/function.py

@@ -38,6 +38,23 @@ def forward(self, output, target, target_weight):
 
         return loss / num_joints
 
+class JointsCELoss(nn.Module):
+    def __init__(self):
+        super(JointsCELoss, self).__init__()
+        self.criterion = nn.MSELoss(reduction='mean').cuda()
+
+    def forward(self, output, target):
+        batch_size = output.size(0)
+        num_joints = output.size(1)
+        loss = 0
+
+        for idx in range(num_joints):
+            class_gt = target[:, idx].view(batch_size)
+            class_pred = output[:,idx]
+            loss += self.criterion(class_pred, class_gt)
+
+        return loss / num_joints
+
 
 class JointsOHKMMSELoss(nn.Module):
     def __init__(self, use_target_weight, topk=8):