Spaces:
Runtime error
Runtime error
Create label.py
Browse files
label.py
ADDED
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from torch.autograd import Variable as V
|
3 |
+
import torchvision.models as models
|
4 |
+
from torchvision import transforms as trn
|
5 |
+
from torch.nn import functional as F
|
6 |
+
import os
|
7 |
+
import numpy as np
|
8 |
+
import cv2
|
9 |
+
from PIL import Image
|
10 |
+
|
11 |
+
|
12 |
+
|
13 |
+
def recursion_change_bn(module):
|
14 |
+
if isinstance(module, torch.nn.BatchNorm2d):
|
15 |
+
module.track_running_stats = 1
|
16 |
+
else:
|
17 |
+
for i, (name, module1) in enumerate(module._modules.items()):
|
18 |
+
module1 = recursion_change_bn(module1)
|
19 |
+
return module
|
20 |
+
|
21 |
+
def load_labels():
|
22 |
+
# prepare all the labels
|
23 |
+
# scene category relevant
|
24 |
+
file_name_category = 'categories_places365.txt'
|
25 |
+
classes = list()
|
26 |
+
with open(file_name_category) as class_file:
|
27 |
+
for line in class_file:
|
28 |
+
classes.append(line.strip().split(' ')[0][3:])
|
29 |
+
classes = tuple(classes)
|
30 |
+
|
31 |
+
# indoor and outdoor relevant
|
32 |
+
file_name_IO = 'IO_places365.txt'
|
33 |
+
with open(file_name_IO) as f:
|
34 |
+
lines = f.readlines()
|
35 |
+
labels_IO = []
|
36 |
+
for line in lines:
|
37 |
+
items = line.rstrip().split()
|
38 |
+
labels_IO.append(int(items[-1]) -1) # 0 is indoor, 1 is outdoor
|
39 |
+
labels_IO = np.array(labels_IO)
|
40 |
+
|
41 |
+
# scene attribute relevant
|
42 |
+
file_name_attribute = 'labels_sunattribute.txt'
|
43 |
+
with open(file_name_attribute) as f:
|
44 |
+
lines = f.readlines()
|
45 |
+
labels_attribute = [item.rstrip() for item in lines]
|
46 |
+
file_name_W = 'W_sceneattribute_wideresnet18.npy'
|
47 |
+
W_attribute = np.load(file_name_W)
|
48 |
+
|
49 |
+
return classes, labels_IO, labels_attribute, W_attribute
|
50 |
+
|
51 |
+
def hook_feature(module, input, output):
|
52 |
+
features_blobs.append(np.squeeze(output.data.cpu().numpy()))
|
53 |
+
|
54 |
+
def returnCAM(feature_conv, weight_softmax, class_idx):
|
55 |
+
# generate the class activation maps upsample to 256x256
|
56 |
+
size_upsample = (256, 256)
|
57 |
+
nc, h, w = feature_conv.shape
|
58 |
+
output_cam = []
|
59 |
+
for idx in class_idx:
|
60 |
+
cam = weight_softmax[class_idx].dot(feature_conv.reshape((nc, h*w)))
|
61 |
+
cam = cam.reshape(h, w)
|
62 |
+
cam = cam - np.min(cam)
|
63 |
+
cam_img = cam / np.max(cam)
|
64 |
+
cam_img = np.uint8(255 * cam_img)
|
65 |
+
output_cam.append(cv2.resize(cam_img, size_upsample))
|
66 |
+
return output_cam
|
67 |
+
|
68 |
+
def returnTF():
|
69 |
+
# load the image transformer
|
70 |
+
tf = trn.Compose([
|
71 |
+
trn.Resize((224,224)),
|
72 |
+
trn.ToTensor(),
|
73 |
+
trn.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
74 |
+
])
|
75 |
+
return tf
|
76 |
+
|
77 |
+
|
78 |
+
def load_model():
|
79 |
+
# this model has a last conv feature map as 14x14
|
80 |
+
|
81 |
+
model_file = 'wideresnet18_places365.pth.tar'
|
82 |
+
import wideresnet
|
83 |
+
model = wideresnet.resnet18(num_classes=365)
|
84 |
+
checkpoint = torch.load(model_file, map_location=lambda storage, loc: storage)
|
85 |
+
state_dict = {str.replace(k,'module.',''): v for k,v in checkpoint['state_dict'].items()}
|
86 |
+
model.load_state_dict(state_dict)
|
87 |
+
|
88 |
+
# hacky way to deal with the upgraded batchnorm2D and avgpool layers...
|
89 |
+
for i, (name, module) in enumerate(model._modules.items()):
|
90 |
+
module = recursion_change_bn(model)
|
91 |
+
model.avgpool = torch.nn.AvgPool2d(kernel_size=14, stride=1, padding=0)
|
92 |
+
|
93 |
+
model.eval()
|
94 |
+
|
95 |
+
# hook the feature extractor
|
96 |
+
features_names = ['layer4','avgpool'] # this is the last conv layer of the resnet
|
97 |
+
for name in features_names:
|
98 |
+
model._modules.get(name).register_forward_hook(hook_feature)
|
99 |
+
return model
|
100 |
+
|
101 |
+
|
102 |
+
def predict_environment(image):
|
103 |
+
# load the labels
|
104 |
+
classes, labels_IO, labels_attribute, W_attribute = load_labels()
|
105 |
+
|
106 |
+
# load the model
|
107 |
+
features_blobs = []
|
108 |
+
model = load_model()
|
109 |
+
|
110 |
+
# load the transformer
|
111 |
+
tf = returnTF() # image transformer
|
112 |
+
|
113 |
+
# get the softmax weight
|
114 |
+
params = list(model.parameters())
|
115 |
+
weight_softmax = params[-2].data.numpy()
|
116 |
+
weight_softmax[weight_softmax<0] = 0
|
117 |
+
|
118 |
+
img = Image.open(image)
|
119 |
+
input_img = V(tf(img).unsqueeze(0))
|
120 |
+
|
121 |
+
# forward pass
|
122 |
+
logit = model.forward(input_img)
|
123 |
+
h_x = F.softmax(logit, 1).data.squeeze()
|
124 |
+
probs, idx = h_x.sort(0, True)
|
125 |
+
probs = probs.numpy()
|
126 |
+
idx = idx.numpy()
|
127 |
+
|
128 |
+
# output the IO prediction
|
129 |
+
io_image = np.mean(labels_IO[idx[:10]]) # vote for the indoor or outdoor
|
130 |
+
env_image = []
|
131 |
+
if io_image < 0.5:
|
132 |
+
env_image.append('Indoor')
|
133 |
+
#print('--TYPE OF ENVIRONMENT: indoor')
|
134 |
+
else:
|
135 |
+
env_image.append('Outdoor')
|
136 |
+
#print('--TYPE OF ENVIRONMENT: outdoor')
|
137 |
+
|
138 |
+
# output the prediction of scene category
|
139 |
+
#print('--SCENE CATEGORIES:')
|
140 |
+
scene_cat=[]
|
141 |
+
for i in range(0, 5):
|
142 |
+
scene_cat.append('{:.3f} -> {}'.format(probs[i], classes[idx[i]]))
|
143 |
+
#print('{:.3f} -> {}'.format(probs[i], classes[idx[i]]))
|
144 |
+
|
145 |
+
return env_image,scene_cat
|