elsamueldev commited on
Commit
b299d3d
1 Parent(s): 509d552

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -6,7 +6,8 @@ import cv2
6
 
7
  login(os.environ["HF_TOKEN"])
8
 
9
- model = from_pretrained_keras("elsamueldev/cats-dogs")
 
10
 
11
 
12
  def preprocess(img: np.array) -> np.array:
@@ -18,16 +19,19 @@ def preprocess(img: np.array) -> np.array:
18
 
19
  return img
20
 
21
- def predict(img: np.array):
22
  img = preprocess(img)
23
 
24
- dog = model.predict(img)[0][0]
 
 
 
25
  cat = 1 - dog
26
 
27
  return {"dog": dog, "cat": cat}
28
 
29
  gr.Interface(
30
  fn=predict,
31
- inputs="image",
32
  outputs="label"
33
  ).launch()
 
6
 
7
  login(os.environ["HF_TOKEN"])
8
 
9
+ modelv1 = from_pretrained_keras("elsamueldev/cats-dogs")
10
+ modelv2 = from_pretrained_keras("elsamueldev/cats-dogs-v2")
11
 
12
 
13
  def preprocess(img: np.array) -> np.array:
 
19
 
20
  return img
21
 
22
+ def predict(img: np.array, model: str):
23
  img = preprocess(img)
24
 
25
+ if model == "v1":
26
+ dog = modelv1.predict(img)[0][0]
27
+ else:
28
+ dog = modelv2.predict(img)[0][0]
29
  cat = 1 - dog
30
 
31
  return {"dog": dog, "cat": cat}
32
 
33
  gr.Interface(
34
  fn=predict,
35
+ inputs=["image", gr.Dropdown(choices=["v1", "v2"], value="v2")],
36
  outputs="label"
37
  ).launch()