INNER CODE UNIT · Python
main
idealo/image-quality-assessment · src/evaluater/predict.py:33
def main(base_model_name, weights_file, image_source, predictions_file, img_format='jpg'):
# load samples
if os.path.isfile(image_source):
image_dir, samples = image_file_to_json(image_source)
else:
image_dir = image_source
samples = image_dir_to_json(image_dir, img_type='jpg')
# build model and load weights
nima = Nima(base_model_name, weights=None)
nima.build()
nima.nima_model.load_weights(weights_file)
# initialize data generator
data_generator = TestDataGenerator(samples, image_dir, 64, 10, nima.preprocessing_function(),
img_format=img_format)
# get predictions