tomaseo2022 commited on
Commit
8b19113
1 Parent(s): b91a567

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import torchvision.transforms as transforms
4
+ from PIL import Image
5
+
6
+ def predict(image):
7
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
8
+ model = torchvision.models.resnet50(pretrained=True).to(device)
9
+ model.eval()
10
+ transform = transforms.Compose([
11
+ transforms.Resize(256),
12
+ transforms.CenterCrop(224),
13
+ transforms.ToTensor(),
14
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
15
+ ])
16
+
17
+ img = Image.fromarray(image.astype('uint8'), 'RGB')
18
+ img = transform(img).unsqueeze(0).to(device)
19
+ with torch.no_grad():
20
+ output = model(img)
21
+ _, predicted = output.max(1)
22
+ return predicted.item()
23
+
24
+ input_image = gr.inputs.Image()
25
+ output_text = gr.outputs.Textbox()
26
+
27
+ gr.Interface(fn=predict, inputs=input_image, outputs=output_text).launch()