lxy1122 commited on
Commit
d8c36a7
·
1 Parent(s): 2b03f29

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from torchvision import transforms
4
+ import gradio as gr
5
+ import os
6
+
7
+
8
+ os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/artist_classes.txt")
9
+
10
+ model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
11
+ checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
12
+ model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
13
+ model.eval()
14
+
15
+ #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
16
+
17
+
18
+ def inference(input_image):
19
+ preprocess = transforms.Compose([
20
+ transforms.Resize(256),
21
+ transforms.CenterCrop(224),
22
+ transforms.ToTensor(),
23
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
24
+ ])
25
+ input_tensor = preprocess(input_image)
26
+ input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
27
+
28
+ # move the input and model to GPU for speed if available
29
+ if torch.cuda.is_available():
30
+ input_batch = input_batch.to('cuda')
31
+ model.to('cuda')
32
+
33
+ with torch.no_grad():
34
+ output = model(input_batch)
35
+ # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
36
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
37
+
38
+ # Read the categories
39
+ with open("artist_classes.txt", "r") as f:
40
+ categories = [s.strip() for s in f.readlines()]
41
+ # Show top categories per image
42
+ top5_prob, top5_catid = torch.topk(probabilities, 6)
43
+ result = {}
44
+ for i in range(top5_prob.size(0)):
45
+ result[categories[top5_catid[i]]] = top5_prob[i].item()
46
+ return result
47
+
48
+ inputs = gr.inputs.Image(type='pil')
49
+ outputs = gr.outputs.Label(type="confidences",num_top_classes=5)
50
+
51
+ title = "MOBILENET V2"
52
+ description = "Gradio demo for MOBILENET V2, Efficient networks optimized for speed and memory, with residual blocks. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
53
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1801.04381'>MobileNetV2: Inverted Residuals and Linear Bottlenecks</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py'>Github Repo</a></p>"
54
+
55
+ #examples = [
56
+ # ['dog.jpg']
57
+ #]
58
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()