pano heatmap network

Author: nuguziii

lib/core/evaluate.py

@@ -28,7 +28,7 @@ def calc_dists(preds, target, normalize):
     return dists
 
 
-def dist_acc(dists, thr=0.5):
+def dist_acc(dists, thr=0.2):
     ''' Return percentage below threshold while ignoring values with a -1 '''
     dist_cal = np.not_equal(dists, -1)
     num_dist_cal = dist_cal.sum()

lib/core/function.py

@@ -175,7 +175,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir,
 
             idx += num_images
 
-            if i % config.PRINT_FREQ == 0:
+            if i % 100 == 0:
                 msg = 'Test: [{0}/{1}]\t' \
                       'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \

lib/core/loss.py

@@ -55,6 +55,25 @@ def forward(self, output, target):
 
         return loss / num_joints
 
+class JointsDistLoss(nn.Module):
+    def __init__(self):
+        super(JointsDistLoss, self).__init__()
+        self.criterion = nn.MSELoss(reduction='mean').cuda()
+
+    def forward(self, output, target):
+        batch = output.size(0)
+        num_joints = output.size(1)
+
+        output = output.reshape(batch, 32, 2)
+        target = target.reshape(batch, 32, 2)
+
+        loss = self.criterion(output[:, :, 0], target[:, :, 0]) + 0.3*self.criterion(output[:, :, 1], target[:, :, 1])
+
+        #diff = [batch, 32]
+        #diff = torch.sqrt((output[:, :, 0] - target[:, :, 0])**2 + (output[:, :, 1] - target[:, :, 1])**2)
+
+        return loss
+
 
 class JointsOHKMMSELoss(nn.Module):
     def __init__(self, use_target_weight, topk=8):

lib/dataset/pano.py

@@ -79,7 +79,7 @@ def __init__(self, cfg, root, image_set, state, transform=None):
         )
 
         print('==> initializing pano {} data.'.format(state))
-        self.annot_path = os.path.join('C:\\Users\CGIP\Desktop\CenterNet\data\pano', state)
+        self.annot_path = os.path.join('C:\\Users\CGIP\Desktop\github\CenterNet\data\pano', state)
         self.pano = pano.PANODataset(self.annot_path)
 
         # load image file names

lib/dataset/panoDataset.py

@@ -25,10 +25,15 @@ def _loadImg(self, file):
 
     def loadData(self, file):
         original_img = self._loadImg(file)
-        img, anns = self._cropImg(original_img, self._loadAnn(file))
+        anns, mask = self._loadAnn(file)
+        img, anns = self._cropImg(original_img, anns)
         img, anns = self._resize(img, anns)
         img = self._normalize(img)
 
+        #mask = self._normalize(self._resize(self._cropImg(mask)))
+        #cv2.imshow('result', mask)
+        #cv2.waitKey(0)
+
         return img, anns
 
     def show_results(self, debugger, original_image, anns, save=True):
@@ -95,18 +100,21 @@ def _normalize(self, img):
         inp = (inp - mean) / std
         return inp
 
-    def _resize(self, img, anns):
+    def _resize(self, img, anns=None):
         img_new = cv2.resize(img, (512, 256), interpolation=cv2.INTER_CUBIC)
         fx = np.size(img, 1) / 512
         fy = np.size(img, 0) / 256
 
+        if anns==None:
+            return img_new
+
         for ipt in range(1, self.num_joints+1):
             anns[ipt]["bbox"] = (anns[ipt]["bbox"]/np.array([fx, fy, fx, fy])).astype(int)
             anns[ipt]["center"] = (anns[ipt]["center"]/np.array([fx, fy])).astype(int)
 
         return img_new, anns
 
-    def _cropImg(self, img, anns):
+    def _cropImg(self, img, anns=None):
         width = np.size(img, 1)
         height = np.size(img, 0)
 
@@ -129,13 +137,16 @@ def _cropImg(self, img, anns):
         diffW = int((width-newW)/2)
         img = img[diffH:diffH+newH,diffW:diffW+newW,:]
 
+        if anns==None:
+            return img
+
         for ipt in range(1, self.num_joints+1):
             anns[ipt]["bbox"] = anns[ipt]["bbox"] - np.array([diffW + self.w_pad, diffH + self.h_pad, diffW + self.w_pad, diffH + self.h_pad])
             anns[ipt]["center"] = anns[ipt]["center"] - np.array([diffW + self.w_pad, diffH + self.h_pad])
 
         return img, anns
 
-    def _loadAnn(self, file, isBinary=False):
+    def _loadAnn(self, file):
         '''
         :param file: pano jpg image path
         :return data: { 1 : {"bbox":[(0,0),(0,0),(0,0),(0,0)], "center":(0,0), "visible":1/0} , 2: {} , ... }
@@ -149,30 +160,46 @@ def _loadAnn(self, file, isBinary=False):
         f = open(file.replace('jpg', 'txt'), 'r')
         txt = f.readlines()
         f.close()
-        for t in txt:
+
+        for idx, t in enumerate(txt):
             dict = {}
             t = t.replace('\n', '').split(', ')
 
             x_coord = []
             y_coord = []
+            coords = np.zeros((8, 2), np.int32)
 
             for i in range(8):
+                coords[i] = [centerX + int(t[2 * i + 3]), centerY + int(t[2 * i + 4])]
                 x_coord.append(centerX + int(t[2 * i + 3]))
                 y_coord.append(centerY + int(t[2 * i + 4]))
 
-            coords = [min(x_coord), min(y_coord), max(x_coord), max(y_coord)]  # minX, minY, maxX, maxY
+            bboxcoords = [min(x_coord), min(y_coord), max(x_coord), max(y_coord)]  # minX, minY, maxX, maxY
 
-            dict["bbox"] = np.array([coords[0], coords[1], coords[2], coords[3]])  # from top-left corner, clockwise
+            dict["bbox"] = np.array([bboxcoords[0], bboxcoords[1], bboxcoords[2], bboxcoords[3]])  # from top-left corner, clockwise
             dict["center"] = np.array([centerX + int(t[27]), centerY + int(t[28])])
 
             if t[1]=='True':
                 dict["visible"] = 1
             else:
                 dict["visible"] = 0
 
+            dict["mask"] = coords
+
             data[teeth_num[t[0]]] = dict
 
-        return data
+        mask = self.getSegMask(np.zeros_like(img), data)
+
+        return data, mask
+
+    def getSegMask(self, img, data):
+        for i in range(1, 33):
+            visible = data[i]["visible"]
+            mask = data[i]["mask"]
+            if visible==1:
+                cv2.fillConvexPoly(img, mask, (255,0,0))
+
+        return img
 
     def showImage(self, img, anns):
         for ann in anns:

lib/models/__init__.py

@@ -8,10 +8,4 @@
 from __future__ import division
 from __future__ import print_function
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import models.pose_resnet
-import models.pose_hrnet
-import models.pose_hrnet_plus
+import models.panoNet

lib/utils/vis.py

@@ -259,7 +259,7 @@ def save_images_landmark(meta, batch_landmark, batch_classification, prefix, i):
 
         cls = batch_classification[0][i]
 
-        if cls>0:
+        if True:
             cv2.putText(image, teeth_num[i+1], (int(joint_lm[0])+4, int(joint_lm[1])+4), cv2.FONT_ITALIC, 0.4,
                         [255, 0, 0], 1)
             cv2.circle(image, (int(joint_lm[0]), int(joint_lm[1])), 2, [255, 0, 0], 2)

tools/test.py

@@ -25,8 +25,8 @@
 import _init_paths
 from config import cfg
 from config import update_config
-from core.loss import JointsMSELoss, JointsCELoss
-from core.function_plus import test
+from core.loss import JointsMSELoss, JointsCELoss, JointsDistLoss
+from core.function import test
 from utils.utils import create_logger
 
 import dataset
@@ -105,7 +105,8 @@ def main():
 
     # classifierLoss = nn.MSELoss(reduction='mean').cuda()
     classifierLoss = JointsCELoss().cuda()
-    lmloss = nn.MSELoss(reduction='mean').cuda()
+    #lmloss = nn.MSELoss(reduction='mean').cuda()
+    lmloss = JointsDistLoss().cuda()
 
     # Data loading code
     test_dataset = eval('dataset.' + cfg.DATASET.DATASET)(
@@ -124,7 +125,7 @@ def main():
     )
 
     # evaluate on validation set
-    test(cfg, test_loader, test_dataset, model, [classifierLoss, lmloss],
+    test(cfg, test_loader, test_dataset, model, [heatmapLoss, lmloss],
              final_output_dir, tb_log_dir)
 
 

tools/train.py

@@ -26,11 +26,11 @@
 import _init_paths
 from config import cfg
 from config import update_config
-from core.loss import JointsMSELoss, JointsCELoss
+from core.loss import JointsMSELoss, JointsCELoss, JointsDistLoss
 #from core.function import train
 #from core.function import validate
-from core.function_plus import train
-from core.function_plus import validate
+from core.function import train
+from core.function import validate
 from utils.utils import get_optimizer
 from utils.utils import save_checkpoint
 from utils.utils import create_logger
@@ -123,7 +123,8 @@ def main():
 
     #classifierLoss = nn.MSELoss(reduction='mean').cuda()
     classifierLoss = JointsCELoss().cuda()
-    lmloss = nn.MSELoss(reduction='mean').cuda()
+    #lmloss = nn.MSELoss(reduction='mean').cuda()
+    lmloss = JointsDistLoss().cuda()
 
     # Data loading code
     train_dataset = eval('dataset.'+cfg.DATASET.DATASET)(
@@ -185,13 +186,13 @@ def main():
         lr_scheduler.step()
 
         # train for one epoch
-        train(cfg, train_loader, model, [classifierLoss, lmloss], optimizer, epoch,
+        train(cfg, train_loader, model, [heatmapLoss, lmloss], optimizer, epoch,
               final_output_dir, tb_log_dir, writer_dict)
 
 
         # evaluate on validation set
         perf_indicator = validate(
-            cfg, valid_loader, valid_dataset, model, [classifierLoss, lmloss],
+            cfg, valid_loader, valid_dataset, model, [heatmapLoss, lmloss],
             final_output_dir, tb_log_dir, writer_dict
         )