class Mask(Masker):
""" Neural network to process face image into a segmentation mask of the face """
def __init__(self, **kwargs):
git_model_id = 8
model_filename = "Nirkin_300_softmax_v1.h5"
super().__init__(git_model_id=git_model_id,
model_filename=model_filename,
**kwargs)
self.name = "VGG Clear"
self.input_size = 300
self.vram = 2944
self.vram_warnings = 1088 # at BS 1. OOMs at higher batch sizes
self.vram_per_batch = 400
self.batchsize = self.config["batch-size"]
def init_model(self):
self.model = KSession(self.name,
self.model_path,
model_kwargs=dict(),
allow_growth=self.config["allow_growth"],
exclude_gpus=self._exclude_gpus)
self.model.load_model()
self.model.append_softmax_activation(layer_index=-1)
placeholder = np.zeros(
(self.batchsize, self.input_size, self.input_size, 3),
dtype="float32")
self.model.predict(placeholder)
def process_input(self, batch):
""" Compile the detected faces for prediction """
input_ = np.array(
[face.feed_face[..., :3] for face in batch["detected_faces"]],
dtype="float32")
batch["feed"] = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :]
logger.trace("feed shape: %s", batch["feed"].shape)
return batch
def predict(self, batch):
""" Run model to get predictions """
predictions = self.model.predict(batch["feed"])
batch["prediction"] = predictions[..., -1]
return batch
def process_output(self, batch):
""" Compile found faces for output """
return batch
class Mask(Masker):
""" Perform transformation to align and get landmarks """
def __init__(self, **kwargs):
git_model_id = 5
model_filename = "Nirkin_500_softmax_v1.h5"
super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
self.name = "VGG Obstructed"
self.input_size = 500
self.vram = 3936
self.vram_warnings = 1088 # at BS 1. OOMs at higher batchsizes
self.vram_per_batch = 304
self.batchsize = self.config["batch-size"]
def init_model(self):
self.model = KSession(self.name, self.model_path, model_kwargs=dict())
self.model.load_model()
self.model.append_softmax_activation(layer_index=-1)
placeholder = np.zeros((self.batchsize, self.input_size, self.input_size, 3),
dtype="float32")
self.model.predict(placeholder)
def process_input(self, batch):
""" Compile the detected faces for prediction """
input_ = [face.feed_face[..., :3] for face in batch["detected_faces"]]
batch["feed"] = input_ - np.mean(input_, axis=(1, 2))[:, None, None, :]
logger.trace("feed shape: %s", batch["feed"].shape)
return batch
def predict(self, batch):
""" Run model to get predictions """
predictions = self.model.predict(batch["feed"])
batch["prediction"] = predictions[..., 0] * -1.0 + 1.0
return batch
def process_output(self, batch):
""" Compile found faces for output """
return batch