Spaces:
Runtime error
Runtime error
File size: 1,385 Bytes
7951498 |
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 |
import argparse, os, json
import numpy as np
from imageio import imread
from PIL import Image
import torch
import torchvision
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
def build_model(model='resnet101', model_stage=3):
cnn = getattr(torchvision.models, model)(pretrained=True)
layers = [
cnn.conv1,
cnn.bn1,
cnn.relu,
cnn.maxpool,
]
for i in range(model_stage):
name = 'layer%d' % (i + 1)
layers.append(getattr(cnn, name))
model = torch.nn.Sequential(*layers)
# model.cuda()
model.eval()
return model
def run_image(img, model):
mean = np.array([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1)
std = np.array([0.229, 0.224, 0.224]).reshape(1, 3, 1, 1)
image = np.concatenate([img], 0).astype(np.float32)
image = (image / 255.0 - mean) / std
image = torch.FloatTensor(image)
image = torch.autograd.Variable(image, volatile=True)
feats = model(image)
feats = feats.data.cpu().clone().numpy()
return feats
def get_img_feat(cnn_model, img, image_height=224, image_width=224):
img_size = (image_height, image_width)
img = np.array(Image.fromarray(np.uint8(img)).resize(img_size))
img = img.transpose(2, 0, 1)[None]
feats = run_image(img, cnn_model)
_, C, H, W = feats.shape
feat_dset = feats.reshape(1, C, H, W)
return feat_dset
|