Multiple model

Dockerfile

@@ -1,4 +1,4 @@
-FROM docker.io/jackrio/bae_repo:v3
+FROM docker.io/jackrio/bae_repo:v4
 
 USER root
 

predict.py

@@ -7,10 +7,10 @@ from models.model_zoo import BoneAgeEstModelZoo
 device = "cpu"
 
 
-def initialize_model():
+def initialize_model(path):
     # Load model
     model = BoneAgeEstModelZoo(branch="gender", pretrained=True, lr=0.001).load_from_checkpoint(
-        "/bae/output/inception_1024_new_data/epoch=13-step=9828.ckpt", map_location=device)
+        path, map_location=device)
     model.model.eval()
     model.classifier.eval()
     model.gender.eval()
@@ -31,8 +31,8 @@ def predict(image, gender):
         is_female = 1
     else:
         is_female = 0
-
-    model = initialize_model()
+    path = "/bae/output/inception_1024_new_data/epoch=13-step=9828.ckpt"
+    model = initialize_model(path)
 
     processed_image = transform(image=np.array(image, dtype=np.uint8))['image']
     processed_image = processed_image.unsqueeze(0)
@@ -43,13 +43,18 @@ def predict(image, gender):
         'image': processed_image,
         'gender': is_female
     }
-    preds = model(scans)
-    return int(preds)
+    preds_1 = model(scans)
+
+    path = "/bae/output/inception_1024/epoch14_inception_1024_kaggle.ckpt"
+    model = initialize_model(path)
+    preds_2 = model(scans)
+
+    return int((preds_1 + preds_2) / 2)
 
 
 def run():
     image_input = gr.inputs.Image(type="pil", label="Input PNG image")
-    gender_input = gr.Checkbox(label="Is Female?", info="Is the scan of a female?", default=True)
+    gender_input = gr.Checkbox(label="Is Female?", info="Is the scan of a female?", default=False)
     output = gr.outputs.Textbox(label="Predicted Age")
 
     BAE = gr.Interface(