breynolds1247 commited on
Commit
3c1be85
1 Parent(s): dc58ac6

add application file

Browse files
Files changed (1) hide show
  1. app.py +78 -0
app.py ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from torchvision import datasets, models, transforms
4
+ import gradio as gr
5
+ import os
6
+ import torch.nn as nn
7
+
8
+
9
+ os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/artist_classes.txt")
10
+ #os.system("wget https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth")
11
+
12
+ #model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
13
+ #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
14
+ #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
15
+ model = models.vgg16()
16
+ num_ftrs = model.classifier[6].in_features
17
+ model.classifier[6] = nn.Linear(num_ftrs, 6)
18
+ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
19
+ #model = model.to(device)
20
+ model.load_state_dict(torch.load('VGG16_weights_May28.pth',map_location=device))
21
+
22
+ #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
23
+
24
+
25
+ def inference(input_image):
26
+ preprocess = transforms.Compose([
27
+ transforms.Resize(260),
28
+ transforms.CenterCrop(224),
29
+ transforms.ToTensor(),
30
+ #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
31
+ ])
32
+ input_tensor = preprocess(input_image)
33
+ input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
34
+
35
+ # move the input and model to GPU for speed if available
36
+ if torch.cuda.is_available():
37
+ input_batch = input_batch.to('cuda')
38
+ model.to('cuda')
39
+ else:
40
+ model.to('cpu')
41
+
42
+ with torch.no_grad():
43
+ output = model(input_batch)
44
+ # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
45
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
46
+
47
+ # Read the categories
48
+ with open("artist_classes.txt", "r") as f:
49
+ categories = [s.strip() for s in f.readlines()]
50
+
51
+ categories = {
52
+ 0:"vanGogh",
53
+ 1:"Monet",
54
+ 2:"Leonardo da Vinci",
55
+ 3:"Rembrandt",
56
+ 4:"Pablo Picasso",
57
+ 5:"Salvador Dali"
58
+ }
59
+
60
+ # Show top categories per image
61
+ top5_prob, top5_catid = torch.topk(probabilities, 6)
62
+ result = {}
63
+ for i in range(top5_prob.size(0)):
64
+ result[categories[top5_catid[i].item()]] = top5_prob[i].item()
65
+ return result
66
+
67
+ inputs = gr.Image(type='pil')
68
+ outputs = gr.Label(type="confidences",num_top_classes=5)
69
+
70
+ title = "Artist Classifier"
71
+ 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."
72
+ 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>"
73
+
74
+ #examples = [
75
+ # ['dog.jpg']
76
+ #]
77
+ #gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()
78
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False).launch()