Update Code to Enable model selection

app.py

@@ -11,7 +11,7 @@ from huggingface_hub import snapshot_download
 import os
 import gradio as gr
 
-MODEL_REPO = 'sokonana/it107model'
+MODEL_REPO = 'sokonana/it107model'  # default model selected
 PATH_TO_LABELS = 'data/label_map.pbtxt'   
 category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
 
@@ -29,25 +29,24 @@ def load_image_into_numpy_array(path):
     return pil_image_as_numpy_array(image)            
 
 def load_model():
-    # wget.download("https://nyp-aicourse.s3-ap-southeast-1.amazonaws.com/pretrained-models/balloon_model.tar.gz")
-    # tarfile.open("balloon_model.tar.gz").extractall()
     model_path = snapshot_download(MODEL_REPO)
 
     model_dir = os.path.join(model_path, 'saved_model')
     detection_model = tf.saved_model.load(model_dir)
     return detection_model    
 
-# samples_folder = 'test_samples
-# image_path = 'test_samples/sample_balloon.jpeg
-# 
 
-def predict(pilimg):
+def predict(pilimg, model):
+
+    if model == 'ResNet':
+        MODEL_REPO = 'sokonana/it107model2'
 
     image_np = pil_image_as_numpy_array(pilimg)
     return predict2(image_np)
 
 def predict2(image_np):
 
+    detection_model = load_model()
     results = detection_model(image_np)
 
     # different object detection models have additional results
@@ -72,7 +71,7 @@ def predict2(image_np):
     
     return result_pil_img
 
-detection_model = load_model()
+# detection_model = load_model()
 # pil_image = Image.open(image_path)
 # image_arr = pil_image_as_numpy_array(pil_image)
 
@@ -80,6 +79,6 @@ detection_model = load_model()
 # predicted_img.save('predicted.jpg')
 
 gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
+             inputs=[gr.Image(type="pil"), gr.Radio(['MobileNet', 'ResNet'],label='Model Selection')],
              outputs=gr.Image(type="pil")
              ).launch(share=True)