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