ditobprasetio commited on
Commit
cde7851
·
1 Parent(s): b0cb56a

add application files

Browse files
Files changed (2) hide show
  1. app.py +47 -0
  2. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torchvision.models as models
3
+ from PIL import Image
4
+ from efficientnet_pytorch import EfficientNet
5
+ import torchvision.transforms as transforms
6
+ import gradio as gr
7
+ from gradio import components
8
+ import numpy as np
9
+
10
+ def predict(image):
11
+ image = Image.fromarray(np.uint8(image)).convert('RGB')
12
+
13
+ model = EfficientNet.from_name('efficientnet-b7', num_classes=2)
14
+ model_weights_path = 'efficientnetb7_tyrequality_classifier.pth'
15
+ model.load_state_dict(torch.load(model_weights_path))
16
+ model.eval()
17
+
18
+ transform = transforms.Compose([
19
+ transforms.Resize(224),
20
+ transforms.CenterCrop(224),
21
+ transforms.ToTensor(),
22
+ transforms.Normalize([0.485, 0.456, 0.406],
23
+ [0.229, 0.224, 0.225])
24
+ ])
25
+
26
+ input_data = transform(image).unsqueeze(0)
27
+
28
+ class_to_label = {0: 'defective', 1: 'good'}
29
+
30
+ with torch.no_grad():
31
+ output = model(input_data)
32
+
33
+ # Get the predicted class label
34
+ _, predicted_class = torch.max(output, 1)
35
+
36
+ probs = torch.nn.functional.softmax(output, dim=1)
37
+ print(probs, "probs")
38
+ conf, _ = torch.max(probs, 1)
39
+
40
+ result = "Tire status is {} with confidence level in {}%".format(class_to_label[predicted_class.item()], conf.item()*100)
41
+ return result
42
+
43
+ iface = gr.Interface(fn=predict,
44
+ inputs=gr.Image(),
45
+ outputs="textbox")
46
+
47
+ iface.launch(share=True)
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ efficientnet_pytorch
4
+ gradio
5
+ numpy