Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def __init__(self, is_target=False, aim=ENCODER_AIM.SPEED):
super().__init__(is_target)
self.model = None
# I think we should make this an enum, something like: speed, balance, accuracy
self.aim = aim
self._pytorch_wrapper = torch.FloatTensor
self._prepared = False
self._scaler = transforms.Scale((224, 224))
# @TODO Magic numbers with no idea left of how they got here, we should at least have the decency of citing some paper that used these as magic numbers :P
self._normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
self._to_tensor = transforms.ToTensor()
pil_logger = logging.getLogger('PIL')
pil_logger.setLevel(logging.ERROR)
def prepare_encoder(self, priming_data):
if self._prepared:
raise Exception('You can only call "prepare_encoder" once for a given encoder.')
if self.model is None:
if self.aim == ENCODER_AIM.SPEED:
self.model = Img2Vec(model='resnet-18')
elif self.aim == ENCODER_AIM.BALANCE:
self.model = Img2Vec(model='resnext-50-small')
elif self.aim == ENCODER_AIM.ACCURACY:
self.model = Img2Vec(model='resnext-50')
else:
self.model = Img2Vec()
self._prepared = True