| from PIL import Image | |
| import numpy as np | |
| from transformers import ViTImageProcessor, TFViTModel | |
| import keras | |
| import argparse | |
| BASE_MODEL = "google/vit-base-patch16-224" | |
| IMAGE_SIZE = 224 | |
| class Inference: | |
| def __init__(self): | |
| self.vit_model = self._load_vit_model() | |
| self.image_preprocessor = self._load_image_preprocessor() | |
| def predict_rotation(self, image_path): | |
| X = self._preprocess(image_path) | |
| y = self.vit_model.predict(X)[0][0] | |
| return y | |
| def _preprocess(self, image_path): | |
| img = Image.open(image_path) | |
| img = img.resize((IMAGE_SIZE, IMAGE_SIZE)) | |
| img = np.array(img) | |
| X_vit = self.image_preprocessor.preprocess(images=[img], return_tensors="pt")["pixel_values"] | |
| return np.array(X_vit) | |
| def _load_image_preprocessor(self): | |
| print("Loading Image Preprocessor") | |
| return ViTImageProcessor.from_pretrained(BASE_MODEL) | |
| def _load_vit_model(self): | |
| print("Loading Model") | |
| vit_base = TFViTModel.from_pretrained(BASE_MODEL) | |
| img_input = keras.layers.Input(shape=(3,IMAGE_SIZE, IMAGE_SIZE)) | |
| x = vit_base.vit(img_input) | |
| y = keras.layers.Dense(1, activation="linear")(x[-1]) | |
| model = keras.Model(inputs=img_input, outputs=y) | |
| print(model.summary()) | |
| print("Loading Weights") | |
| model.load_weights("weights.h5") | |
| return model | |
| if __name__=="__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--image-path", type=str, required=True) | |
| args = parser.parse_args() | |
| model = Inference() | |
| expected_angle = model.predict_rotation(args.image_path) | |
| print(f"Predicted angle is about '{expected_angle}' degrees") | |