Diego Carpintero commited on
Commit
d68a05b
1 Parent(s): 453f970

add class and object prediction

Browse files
Files changed (1) hide show
  1. app.py +33 -13
app.py CHANGED
@@ -11,21 +11,41 @@ inputs = gr.components.Image()
11
  outputs = gr.components.Label()
12
  examples = "examples"
13
 
14
- model = load_learner("model/model.pkl")
15
- labels = model.dls.vocab
16
 
 
 
17
 
18
- def predict(img):
19
- pred, pred_idx, probs = model.predict(img)
20
- return dict(zip(labels, map(float, probs)))
21
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
22
 
23
- demo = gr.Interface(
24
- fn=predict,
25
- inputs=inputs,
26
- outputs=outputs,
27
- examples=examples,
28
- title=title,
29
- description=description,
30
- ).queue(default_concurrency_limit=5)
31
  demo.launch()
 
11
  outputs = gr.components.Label()
12
  examples = "examples"
13
 
14
+ model_class = load_learner("models/model.class.pkl")
15
+ labels_class = model_class.dls.vocab
16
 
17
+ model_object = load_learner("models/model.object.pkl")
18
+ labels_object = model_object.dls.vocab
19
 
 
 
 
20
 
21
+ def predict_class(img):
22
+ pred, pred_idx, probs = model_class.predict(img)
23
+ return dict(zip(labels_class, map(float, probs)))
24
+
25
+
26
+ def predict_object(img):
27
+ pred, pred_idx, probs = model_object.predict(img)
28
+ return dict(zip(labels_object, map(float, probs)))
29
+
30
+
31
+ with gr.Blocks() as demo:
32
+ with gr.Tab("Class Prediction"):
33
+ gr.Interface(
34
+ fn=predict_class,
35
+ inputs=inputs,
36
+ outputs=outputs,
37
+ examples=examples,
38
+ title=title,
39
+ description=description,
40
+ ).queue(default_concurrency_limit=5)
41
+ with gr.Tab("Object Prediction"):
42
+ gr.Interface(
43
+ fn=predict_object,
44
+ inputs=inputs,
45
+ outputs=outputs,
46
+ examples=examples,
47
+ title=title,
48
+ description=description,
49
+ ).queue(default_concurrency_limit=5)
50
 
 
 
 
 
 
 
 
 
51
  demo.launch()