def __init__(self, pretrained=True, min_size=10, verbose=True, **kwargs):
self.verbose = verbose
pnet = p_net()
self.rnet = ImageDetectionModel(input_shape=(3, 24, 24),
output=r_net())._model
self.onet = ImageDetectionModel(input_shape=(3, 48, 48),
output=o_net())._model
if pretrained == True:
pnet = Pnet()._model
self.rnet = Rnet()._model
self.onet = Onet()._model
self.min_size = min_size
super(Mtcnn, self).__init__(input_shape=(3, 224, 224), output=pnet)
self.pnet = pnet
self.signature = get_signature(self._model.forward)
#data preprocess
self.preprocess_flow = [normalize(0, 255)]
self.nms_threshould = [0.9, 0.9, 0.3]
self.detection_threshould = [0.5, 0.6, 0.9]
pnet.add_module(
'pnet_detector',
DetectorHead(cellsize=12, threshould=0.5, min_size=self.min_size))
def Onet(pretrained=True, input_shape=(3, 48, 48), **kwargs):
if input_shape is not None and len(input_shape) == 3:
input_shape = tuple(input_shape)
else:
input_shape = (3, 48, 48)
onet = ImageDetectionModel(input_shape=(3, 48, 48), output=o_net())
onet.preprocess_flow = [Normalize(0, 255), image_backend_adaption]
if pretrained == True:
download_model_from_google_drive('1a1dAlSzJOAfIz77Ic38JMQJYWDG_b7-_',
dirname, 'onet.pth')
recovery_model = load(os.path.join(dirname, 'onet.pth'))
recovery_model = fix_layer(recovery_model)
recovery_model.to(_device)
onet.model = recovery_model
return onet
def Rnet(pretrained=True, input_shape=(3, 24, 24), **kwargs):
if input_shape is not None and len(input_shape) == 3:
input_shape = tuple(input_shape)
else:
input_shape = (3, 24, 24)
rnet = ImageDetectionModel(input_shape=(3, 24, 24), output=r_net())
rnet.preprocess_flow = [Normalize(0, 255), image_backend_adaption]
if pretrained == True:
download_model_from_google_drive('1CH7z133_KrcWMx9zXAblMCV8luiQ3wph',
dirname, 'rnet.pth')
recovery_model = load(os.path.join(dirname, 'rnet.pth'))
recovery_model = fix_layer(recovery_model)
recovery_model.to(_device)
rnet.model = recovery_model
return rnet
def Pnet(pretrained=True, input_shape=(3, 12, 12), **kwargs):
if input_shape is not None and len(input_shape) == 3:
input_shape = tuple(input_shape)
else:
input_shape = (3, 12, 12)
pnet = ImageDetectionModel(input_shape=(3, 12, 12), output=p_net())
pnet.preprocess_flow = [normalize(0, 255), image_backend_adaption]
if pretrained == True:
download_model_from_google_drive('1w9ahipO8D9U1dAXMc2BewuL0UqIBYWSX',
dirname, 'pnet.pth')
recovery_model = torch.load(os.path.join(dirname, 'pnet.pth'))
recovery_model = fix_layer(recovery_model)
recovery_model.to(_device)
pnet.model = recovery_model
return pnet
def Pnet(pretrained=True,
input_shape=(3, 12, 12),
freeze_features=True,
**kwargs):
if input_shape is not None and len(input_shape) == 3:
input_shape = tuple(input_shape)
else:
input_shape = (3, 224, 224)
pnet = ImageDetectionModel(input_shape=(3, 12, 12), output=p_net())
if pretrained == True:
download_model_from_google_drive('1w9ahipO8D9U1dAXMc2BewuL0UqIBYWSX',
dirname, 'pnet.pth')
recovery_model = fix_layer(load(os.path.join(dirname, 'pnet.pth')))
pnet.model = recovery_model
pnet.model.input_shape = input_shape
pnet.model.to(_device)
return pnet