def determine_image(prediction_model, input_path, logger, image_name):
filepath = join(input_path, image_name)
# Check image extension (it should be JPEG or PNG)
try:
check_mimetype(filepath)
except WrongMimeTypeError:
logger.info(
f"{'*' * 50}\n{basename(image_name)} - unsupported_file\n{'*' * 50}"
)
try:
image = Image.open(filepath)
result = prediction_model.predict(image)
logger.info(f"{'*' * 50}\n{image_name} - {result}\n{'*' * 50}")
except PIL.UnidentifiedImageError as e:
logger.info(
f"{'*' * 50}\n{basename(image_name)} - cannot identify your image\n{'*' * 50}"
)
def predict(prediction_model):
predict_model = get_prediction_model(prediction_model)
if predict_model:
image = flask.request.files["image"]
if flask.request.files.get("image"):
try:
# Check image extension (it should be JPEG or PNG)
check_mimetype(flask.request.files.get("image"))
except WrongMimeTypeError:
return abort(400, "Unsupported file")
image.seek(0)
image = image.read()
image = Image.open(io.BytesIO(image))
result = predict_model().predict(image)
resp = make_response({"data": result}, 200)
return resp
return abort(400, "There is no image to process")
return abort(400, f"Prediction model with the name of the '{prediction_model}' does not exist")
def execute(predict_model, input_path):
prediction_model = get_prediction_model(predict_model)
if not prediction_model:
raise Exception(
f"Prediction model with the name of the '{predict_model}' does not exist"
)
if not exists(input_path):
raise Exception(f"Images path does not exist")
logger = init_logger(predict_model)
prediction_model = prediction_model()
if isdir(input_path):
# Get all files from directory
images = [
f for f in os.listdir(input_path) if isfile(join(input_path, f))
]
images.sort()
for img in tqdm.tqdm(images, total=len(images)):
determine_image(prediction_model, input_path, logger, img)
else:
try:
check_mimetype(input_path)
image = Image.open(input_path)
except WrongMimeTypeError as e:
logger.info(
f"{'*' * 50}\n{basename(input_path)} - unsupported_file\n{'*' * 50}"
)
raise e
except PIL.UnidentifiedImageError as e:
logger.info(
f"{'*' * 50}\n{basename(input_path)} - cannot identify your image\n{'*' * 50}"
)
raise e
result = prediction_model.predict(image)
logger.info(
f"{'*' * 50}\n{basename(input_path)} - {result}\n{'*' * 50}")
def test_check_mimetype_with_unsupported_type(self):
image = "no_jpeg_no_png.svg"
image_path = self.fixtures_dir + image
with self.assertRaises(WrongMimeTypeError):
check_mimetype(image_path)
def test_check_mimetype_using_image_object(self):
for image, result in FIXTURES.items():
with open(self.fixtures_dir + image, "rb") as image:
mime_type = check_mimetype(image)
self.assertEqual(mime_type, result)
def test_check_mimetype_using_filepath(self):
for image, result in FIXTURES.items():
image_path = self.fixtures_dir + image
mime_type = check_mimetype(image_path)
self.assertEqual(mime_type, result)