mirari commited on
Commit
603d751
·
1 Parent(s): a685e4d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -4
app.py CHANGED
@@ -10,8 +10,15 @@ 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])])
@@ -19,12 +26,14 @@ def transform_image(image):
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
@@ -32,6 +41,7 @@ def predict(img):
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)
 
 
 
10
  import random
11
 
12
  import torchvision.transforms as transforms
13
+ import PIL
14
+
15
+ import gradio as gr
16
+
17
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18
 
19
  model = load_learner('export (2).pkl')
20
+ model.cpu()
21
+ model.eval()
22
 
23
  def transform_image(image):
24
  my_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalzie([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
 
26
 
27
  def predict(img):
28
  img = PILImage.create(img)
29
+
30
  image = transforms.Resize((480,640))(img)
31
  tensor = transform_image(image=image)
 
32
  with torch.no_grad():
33
  outputs = model(tensor)
34
+
35
+ outputs = torch.argmax(outputs,1)
36
+
37
  mask = np.array(outputs.cpu())
38
  mask[mask==0]=255
39
  mask[mask==1]=150
 
41
  mask[mask==3]=25
42
  mask[mask==4]=0
43
  mask=np.reshape(mask,(480,640))
 
44
 
45
+ return Image.fromarray(mask.astype('uint8'))
46
+
47
+ gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128,128)), outputs=gr.inputs.Image(), examples=['color_157.jpg','color_158.jpg']).launch(share=False)