Merge branch 'master' of https://github.com/deepinsight/insightface

RetinaFace/rcnn/dataset/retinaface.py

@@ -7,7 +7,7 @@ import cv2
 import os
 import numpy as np
 import json
-from PIL import Image
+#from PIL import Image
 
 from ..logger import logger
 from .imdb import IMDB
@@ -64,7 +64,11 @@ class retinaface(IMDB):
         nonattr_box_num = 0
         landmark_num = 0
 
+        pp = 0
         for fp in self._fp_bbox_map:
+            pp += 1
+            if pp%1000==0:
+              print('loading', pp)
             if self._split=='test':
               image_path = os.path.join(self._imgs_path, fp)
               roi = {'image': image_path}
@@ -78,10 +82,11 @@ class retinaface(IMDB):
             gt_classes = np.ones([len(self._fp_bbox_map[fp])], np.int32)
             overlaps = np.zeros([len(self._fp_bbox_map[fp]), 2], np.float)
 
+            imsize = cv2.imread(os.path.join(self._imgs_path, fp)).shape[0:2][::-1]
             ix = 0
 
             for aline in self._fp_bbox_map[fp]:
-                imsize = Image.open(os.path.join(self._imgs_path, fp)).size
+                #imsize = Image.open(os.path.join(self._imgs_path, fp)).size
                 values = [float(x) for x in aline.strip().split()]
                 bbox = [values[0], values[1], values[0]+values[2], values[1]+values[3]]
 
@@ -100,7 +105,7 @@ class retinaface(IMDB):
                     continue
 
                 boxes[ix, :] = np.array([x1, y1, x2, y2], np.float)
-                if self._split.startswith('train'):
+                if self._split=='train':
                   landmark = np.array( values[4:19], dtype=np.float32 ).reshape((5,3))
                   for li in range(5):
                     #print(landmark)
@@ -129,7 +134,7 @@ class retinaface(IMDB):
                 ix += 1
             max_num_boxes = max(max_num_boxes, ix)
             #overlaps = scipy.sparse.csr_matrix(overlaps)
-            if self._split.startswith('train') and ix==0:
+            if self._split=='train' and ix==0:
               continue
             boxes = boxes[:ix,:]
             landmarks = landmarks[:ix,:,:]

RetinaFace/rcnn/sample_config.py

@@ -66,6 +66,7 @@ config.CONTEXT_FILTER_RATIO = 1
 config.max_feat_channel = 9999
 
 config.USE_CROP = True
+config.USE_FPN = True
 config.USE_DCN = 0
 config.FACE_LANDMARK = True
 config.USE_OCCLUSION = False

RetinaFace/rcnn/symbol/symbol_common.py

@@ -236,6 +236,21 @@ def get_sym_conv(data, sym):
     c1 = stride2layer[8]
     c2 = stride2layer[16]
     c3 = stride2layer[32]
+    if not config.USE_FPN:
+      assert len(config.RPN_ANCHOR_CFG)==3
+      c3 = conv_act_layer(c3, 'rf_c3_lateral',
+          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
+      c2_lateral = conv_act_layer(c2, 'rf_c2_lateral',
+          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
+      c2 = c2_lateral
+      c1_lateral = conv_act_layer(c1, 'rf_c1_red_conv',
+          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
+      c1 = c1_lateral
+      ret = {8: c1, 16:c2, 32: c3}
+      return ret
+
+
+
     c3 = conv_act_layer(c3, 'rf_c3_lateral',
         F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
     #c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)

RetinaFace/train.py

@@ -118,7 +118,8 @@ def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch,
     logger.info('output shape %s' % pprint.pformat(out_shape_dict))
 
 
-    for k,v in arg_shape_dict.iteritems():
+    for k in arg_shape_dict:
+      v = arg_shape_dict[k]
       if k.find('upsampling')>=0:
         print('initializing upsampling_weight', k)
         arg_params[k] = mx.nd.zeros(shape=v)
@@ -201,6 +202,7 @@ def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch,
     lr_epoch = [int(epoch) for epoch in lr_step.split(',')]
     lr_epoch_diff = [epoch - begin_epoch for epoch in lr_epoch if epoch > begin_epoch]
     lr_iters = [int(epoch * len(roidb) / input_batch_size) for epoch in lr_epoch_diff]
+    iter_per_epoch = int(len(roidb)/input_batch_size)
 
     lr_steps = []
     if len(lr_iters)==5:
@@ -300,6 +302,9 @@ def train_net(args, ctx, pretrained, epoch, prefix, begin_epoch, end_epoch,
           print('lr change to', opt.lr,' in batch', mbatch, file=sys.stderr)
           break
 
+      if mbatch%iter_per_epoch==0:
+        print('saving checkpoint', mbatch, file=sys.stderr)
+        save_model(0)
       if mbatch==lr_steps[-1][0]:
         print('saving final checkpoint', mbatch, file=sys.stderr)
         save_model(0)