Update infer-rotation.py
Browse files- infer-rotation.py +1 -2
infer-rotation.py
CHANGED
|
@@ -5,7 +5,6 @@ from transformers import ViTImageProcessor, TFViTModel
|
|
| 5 |
import keras
|
| 6 |
import argparse
|
| 7 |
|
| 8 |
-
VIT_WEIGHTS_PATH = "model-vit-ang-loss.h5"
|
| 9 |
BASE_MODEL = "google/vit-base-patch16-224"
|
| 10 |
IMAGE_SIZE = 224
|
| 11 |
|
|
@@ -43,7 +42,7 @@ class Inference:
|
|
| 43 |
print(model.summary())
|
| 44 |
|
| 45 |
print("Loading Weights")
|
| 46 |
-
model.load_weights(
|
| 47 |
|
| 48 |
return model
|
| 49 |
|
|
|
|
| 5 |
import keras
|
| 6 |
import argparse
|
| 7 |
|
|
|
|
| 8 |
BASE_MODEL = "google/vit-base-patch16-224"
|
| 9 |
IMAGE_SIZE = 224
|
| 10 |
|
|
|
|
| 42 |
print(model.summary())
|
| 43 |
|
| 44 |
print("Loading Weights")
|
| 45 |
+
model.load_weights("weights.h5")
|
| 46 |
|
| 47 |
return model
|
| 48 |
|