valiyevfagan commited on
Commit
a65b579
·
verified ·
1 Parent(s): cc59fe8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -6
app.py CHANGED
@@ -1,20 +1,18 @@
1
  from fastai.vision.all import *
2
  import gradio as gr
3
 
4
- # Example custom label function used during training
5
- def get_label(fn): return fn.parent.name
6
-
7
- # If any custom transform class or model wrapper was used, define/import it here
8
-
9
  # Load the exported model
10
- learn = load_learner('trash_model(1).pkl')
11
 
 
12
  labels = learn.dls.vocab
13
 
 
14
  def classify_trash(img):
15
  pred_class, pred_idx, probs = learn.predict(img)
16
  return {labels[i]: float(probs[i]) for i in range(len(labels))}
17
 
 
18
  interface = gr.Interface(
19
  fn=classify_trash,
20
  inputs=gr.Image(type="pil"),
@@ -23,4 +21,5 @@ interface = gr.Interface(
23
  description="Upload a trash image to classify it into one of 5 categories."
24
  )
25
 
 
26
  interface.launch()
 
1
  from fastai.vision.all import *
2
  import gradio as gr
3
 
 
 
 
 
 
4
  # Load the exported model
5
+ learn = load_learner('trash_model(1s).pkl')
6
 
7
+ # Define labels (make sure they match your model's training labels)
8
  labels = learn.dls.vocab
9
 
10
+ # Define prediction function
11
  def classify_trash(img):
12
  pred_class, pred_idx, probs = learn.predict(img)
13
  return {labels[i]: float(probs[i]) for i in range(len(labels))}
14
 
15
+ # Gradio Interface
16
  interface = gr.Interface(
17
  fn=classify_trash,
18
  inputs=gr.Image(type="pil"),
 
21
  description="Upload a trash image to classify it into one of 5 categories."
22
  )
23
 
24
+ # Launch app
25
  interface.launch()