Commit d9514269 authored by zhongshangwu's avatar zhongshangwu

适配Python3

parent 3c86aeb3
...@@ -6,8 +6,10 @@ import math ...@@ -6,8 +6,10 @@ import math
import cv2 import cv2
from multiprocessing import Pool from multiprocessing import Pool
from itertools import repeat from itertools import repeat
#from itertools import zip_longest as zip try:
from itertools import izip from itertools import izip
except ImportError:
izip = zip
from helper import nms, adjust_input, generate_bbox, detect_first_stage_warpper from helper import nms, adjust_input, generate_bbox, detect_first_stage_warpper
class MtcnnDetector(object): class MtcnnDetector(object):
...@@ -49,7 +51,7 @@ class MtcnnDetector(object): ...@@ -49,7 +51,7 @@ class MtcnnDetector(object):
# load 4 models from folder # load 4 models from folder
models = ['det1', 'det2', 'det3','det4'] models = ['det1', 'det2', 'det3','det4']
models = [ os.path.join(model_folder, f) for f in models] models = [ os.path.join(model_folder, f) for f in models]
self.PNets = [] self.PNets = []
for i in range(num_worker): for i in range(num_worker):
workner_net = mx.model.FeedForward.load(models[0], 1, ctx=ctx) workner_net = mx.model.FeedForward.load(models[0], 1, ctx=ctx)
...@@ -115,7 +117,7 @@ class MtcnnDetector(object): ...@@ -115,7 +117,7 @@ class MtcnnDetector(object):
bbox[:, 0:4] = bbox[:, 0:4] + aug bbox[:, 0:4] = bbox[:, 0:4] + aug
return bbox return bbox
def pad(self, bboxes, w, h): def pad(self, bboxes, w, h):
""" """
pad the the bboxes, alse restrict the size of it pad the the bboxes, alse restrict the size of it
...@@ -185,7 +187,7 @@ class MtcnnDetector(object): ...@@ -185,7 +187,7 @@ class MtcnnDetector(object):
yield l[i:i + n] yield l[i:i + n]
num_list = range(number) num_list = range(number)
return list(chunks(num_list, self.num_worker)) return list(chunks(num_list, self.num_worker))
def detect_face_limited(self, img, det_type=2): def detect_face_limited(self, img, det_type=2):
height, width, _ = img.shape height, width, _ = img.shape
if det_type>=2: if det_type>=2:
...@@ -257,7 +259,7 @@ class MtcnnDetector(object): ...@@ -257,7 +259,7 @@ class MtcnnDetector(object):
pick = nms(total_boxes, 0.7, 'Min') pick = nms(total_boxes, 0.7, 'Min')
total_boxes = total_boxes[pick] total_boxes = total_boxes[pick]
points = points[pick] points = points[pick]
if not self.accurate_landmark: if not self.accurate_landmark:
return total_boxes, points return total_boxes, points
...@@ -350,7 +352,7 @@ class MtcnnDetector(object): ...@@ -350,7 +352,7 @@ class MtcnnDetector(object):
# return_boxes = self.detect_first_stage(img, scale, 0) # return_boxes = self.detect_first_stage(img, scale, 0)
# if return_boxes is not None: # if return_boxes is not None:
# total_boxes.append(return_boxes) # total_boxes.append(return_boxes)
sliced_index = self.slice_index(len(scales)) sliced_index = self.slice_index(len(scales))
total_boxes = [] total_boxes = []
for batch in sliced_index: for batch in sliced_index:
...@@ -359,14 +361,14 @@ class MtcnnDetector(object): ...@@ -359,14 +361,14 @@ class MtcnnDetector(object):
local_boxes = map( detect_first_stage_warpper, \ local_boxes = map( detect_first_stage_warpper, \
izip(repeat(img), self.PNets[:len(batch)], [scales[i] for i in batch], repeat(self.threshold[0])) ) izip(repeat(img), self.PNets[:len(batch)], [scales[i] for i in batch], repeat(self.threshold[0])) )
total_boxes.extend(local_boxes) total_boxes.extend(local_boxes)
#print('local_boxes',local_boxes) #print('local_boxes',local_boxes)
# remove the Nones # remove the Nones
total_boxes = [ i for i in total_boxes if i is not None] total_boxes = [ i for i in total_boxes if i is not None]
if len(total_boxes) == 0: if len(total_boxes) == 0:
return None return None
total_boxes = np.vstack(total_boxes) total_boxes = np.vstack(total_boxes)
if total_boxes.size == 0: if total_boxes.size == 0:
...@@ -466,7 +468,7 @@ class MtcnnDetector(object): ...@@ -466,7 +468,7 @@ class MtcnnDetector(object):
pick = nms(total_boxes, 0.7, 'Min') pick = nms(total_boxes, 0.7, 'Min')
total_boxes = total_boxes[pick] total_boxes = total_boxes[pick]
points = points[pick] points = points[pick]
if not self.accurate_landmark: if not self.accurate_landmark:
return total_boxes, points return total_boxes, points
...@@ -521,7 +523,7 @@ class MtcnnDetector(object): ...@@ -521,7 +523,7 @@ class MtcnnDetector(object):
input list input list
Retures: Retures:
------- -------
colMat: colMat:
""" """
assert len(pts_list) > 0 assert len(pts_list) > 0
...@@ -537,8 +539,8 @@ class MtcnnDetector(object): ...@@ -537,8 +539,8 @@ class MtcnnDetector(object):
find transform between shapes find transform between shapes
Parameters: Parameters:
---------- ----------
from_shape: from_shape:
to_shape: to_shape:
Retures: Retures:
------- -------
tran_m: tran_m:
...@@ -599,7 +601,7 @@ class MtcnnDetector(object): ...@@ -599,7 +601,7 @@ class MtcnnDetector(object):
Retures: Retures:
------- -------
crop_imgs: list, n crop_imgs: list, n
cropped and aligned faces cropped and aligned faces
""" """
crop_imgs = [] crop_imgs = []
for p in points: for p in points:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment