yuechen-yang commited on
Commit
718e70b
·
1 Parent(s): 433d298

Add application file

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+
4
+ from transformers import pipeline
5
+ import numpy as np
6
+ from transformers import AutoFeatureExtractor, AutoModelForImageClassification
7
+
8
+
9
+ st.set_page_config(layout='wide',
10
+ page_title='Garbage image classification'
11
+ )
12
+
13
+
14
+
15
+
16
+ def main():
17
+
18
+
19
+
20
+ st.header("Try it out!")
21
+
22
+ uploaded_file = st.file_uploader("Upload Files",type=['png','jpeg'])
23
+ img=Image.open(uploaded_file)
24
+
25
+ extractor = AutoFeatureExtractor.from_pretrained("yangy50/garbage-classification")
26
+ model = AutoModelForImageClassification.from_pretrained("yangy50/garbage-classification")
27
+
28
+ inputs = extractor(img,return_tensors="pt")
29
+ outputs = model(**inputs)
30
+ label_num=outputs.logits.softmax(1).argmax(1)
31
+ label_num=label_num.item()
32
+
33
+ if label_num==0:
34
+ st.write("cardboard")
35
+ elif label_num==1:
36
+ st.write("glass")
37
+ elif label_num==2:
38
+ st.write("metal")
39
+ elif label_num==3:
40
+ st.write("paper")
41
+ elif label_num==4:
42
+ st.write("plastic")
43
+ else:
44
+ st.write("trash")
45
+
46
+
47
+
48
+ if __name__ == '__main__':
49
+ main()