mirari commited on
Commit
e8bd95a
·
1 Parent(s): 8bf0f04

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastai.basics import *
2
+ from fastai.vision import models
3
+ from fastai.vision.all import *
4
+ from fastai.metrics import *
5
+ from fastai.data.all import *
6
+ from fastai.callback import *
7
+
8
+
9
+ from pathlib import Path
10
+ import random
11
+
12
+ import torchvision.transforms as transforms
13
+
14
+ model = load_learner('export (2).pkl')
15
+
16
+ def transform_image(image):
17
+ my_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalzie([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
18
+ return my_transforms(image).unsqueeze(0).to(device)
19
+
20
+ def predict(img):
21
+ img = PILImage.create(img)
22
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
+ image = transforms.Resize((480,640))(img)
24
+ tensor = transform_image(image=image)
25
+ model.to(device)
26
+ with torch.no_grad():
27
+ outputs = model(tensor)
28
+ mask = np.array(outputs.cpu())
29
+ mask[mask==0]=255
30
+ mask[mask==1]=150
31
+ mask[mask==2]=76
32
+ mask[mask==3]=25
33
+ mask[mask==4]=0
34
+ mask=np.reshape(mask,(480,640))
35
+ Image.fromarray(mask.astype('uint8'))
36
+
37
+ gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128,128)), outputs=gr.inputs.Image(shape=(128,128)), examples=['color_157.jpg','color_158.jpg']).launch(share=False)