Update app.py

app.py

@@ -6,7 +6,8 @@ import cv2
 
 login(os.environ["HF_TOKEN"])
 
-model = from_pretrained_keras("elsamueldev/cats-dogs")
+modelv1 = from_pretrained_keras("elsamueldev/cats-dogs")
+modelv2 = from_pretrained_keras("elsamueldev/cats-dogs-v2")
 
 
 def preprocess(img: np.array) -> np.array:
@@ -18,16 +19,19 @@ def preprocess(img: np.array) -> np.array:
 
     return img
 
-def predict(img: np.array):
+def predict(img: np.array, model: str):
     img = preprocess(img)
 
-    dog = model.predict(img)[0][0]
+    if model == "v1":
+        dog = modelv1.predict(img)[0][0]
+    else:
+        dog = modelv2.predict(img)[0][0]
     cat = 1 - dog
 
     return {"dog": dog, "cat": cat}
 
 gr.Interface(
     fn=predict,
-    inputs="image",
+    inputs=["image", gr.Dropdown(choices=["v1", "v2"], value="v2")],
     outputs="label"
 ).launch()