luis56125 commited on
Commit
03c8d09
1 Parent(s): d9ac45f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import from_pretrained_fastai
2
+ import gradio as gr
3
+ from fastai.vision.all import *
4
+
5
+ from pathlib import Path
6
+ import PIL
7
+ import torchvision.transforms as transforms
8
+
9
+
10
+
11
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
+ model = torch.jit.load("unet.pth")
13
+ model = model.cpu()
14
+ model.eval()
15
+
16
+ def transform_image(image):
17
+ my_transforms = transforms.Compose([transforms.ToTensor(),
18
+ transforms.Normalize(
19
+ [0.485, 0.456, 0.406],
20
+ [0.229, 0.224, 0.225])])
21
+ image_aux = image
22
+
23
+ image = transforms.Resize((480,640))(Image.fromarray(image))
24
+ tensor = my_transforms(image_aux).unsqueeze(0).to(device)
25
+
26
+
27
+ model.to(device)
28
+ with torch.no_grad():
29
+ outputs = model(tensor)
30
+
31
+ outputs = torch.argmax(outputs,1)
32
+
33
+
34
+ mask = np.array(outputs.cpu())
35
+ mask[mask==0]=255
36
+ mask[mask==1]=150
37
+ mask[mask==2]=76
38
+ mask[mask==3]=25
39
+ mask[mask==4]=0
40
+
41
+ mask=np.reshape(mask,(480,640))
42
+ return Image.fromarray(mask.astype('uint8'))
43
+
44
+
45
+
46
+ # Creamos la interfaz y la lanzamos.
47
+ gr.Interface(fn=transform_image, inputs=gr.inputs.Image(shape=(640, 480)), outputs=gr.outputs.Image(),examples=['color_188.jpg','color_189.jpg']).launch(share=False)
48
+
49
+