forked from Tencent/tencent-ml-images
-
Notifications
You must be signed in to change notification settings - Fork 0
/
extract_feature.py
105 lines (96 loc) · 3.54 KB
/
extract_feature.py
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
#!/usr/bin/python
"""
Tencent is pleased to support the open source community by making Tencent ML-Images available.
Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved.
Licensed under the BSD 3-Clause License (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
https://opensource.org/licenses/BSD-3-Clause
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
"""
"""Use pre-trained model extract image feature
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import numpy as np
import cv2 as cv
import tensorflow as tf
from models import resnet as resnet
from flags import FLAGS
tf.app.flags.DEFINE_string("result", "",
"file name to save features")
tf.app.flags.DEFINE_string("images", "",
"contains image path per line per image")
"""Crop Image To 224*224
Args:
img: an 3-D numpy array (H,W,C)
type: crop method support [ center | 10crop ]
"""
def preprocess(img, type="center"):
# resize image with smallest side to be 256
rawH = float(img.shape[0])
rawW = float(img.shape[1])
newH = 256.0
newW = 256.0
if rawH <= rawW:
newW = (rawW/rawH) * newH
else:
newH = (rawH/rawW) * newW
img = cv.resize(img, (int(newW), int(newH)))
imgs = None
if type=='center':
imgs = np.zeros((1, 224, 224, 3))
imgs[0,...] = img[int((newH-224)/2):int((newH-224)/2)+224,
int((newW-224)/2):int((newW-224)/2)+224]
elif type=='10crop':
imgs = np.zeros((10, 224, 224, 3))
offset = [(0, 0),
(0, int(newW-224)),
(int(newH-224), 0),
(int(newH-224), int(newW-224)),
(int((newH-224)/2), int((newW-224)/2))]
for i in range(0, 5):
imgs[i,...] = img[offset[i][0]:offset[i][0]+224,
offset[i][1]:offset[i][1]+224]
img = cv.flip(img, 1)
for i in range(0, 5):
imgs[i+5,...] = img[offset[i][0]:offset[i][0]+224,
offset[i][1]:offset[i][1]+224]
else:
raise ValueError("Type not support")
imgs = ((imgs/255.0) - 0.5) * 2.0
imgs = imgs[...,::-1]
return imgs
# build model
images = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3])
net = resnet.ResNet(images, is_training=False)
net.build_model()
logits = net.logit
feat = net.feat
# restore model
saver = tf.train.Saver(tf.global_variables())
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = str(FLAGS.visiable_gpu)
config.log_device_placement=False
sess = tf.Session(config=config)
# load trained model
saver.restore(sess, FLAGS.pretrain_ckpt)
# inference on net
types='center'
ffeat = open(FLAGS.result, 'w')
with open(FLAGS.images, 'r') as lines:
for line in lines:
sp = line.rstrip('\n').split(' ')
raw_img = cv.imread(sp[0])
if type(raw_img)==None or raw_img.data==None :
print("open pic " + sp[0] + " failed")
continue
imgs = preprocess(raw_img, types)
feats = sess.run(feat, {images:imgs})
feats = np.squeeze(feats[0])
if types=='10crop':
feats = np.mean(feats, axis=0)
print('feature-length:{}, feature={}'.format(len(feats), feats))
ffeat.write(sp[0] + "\t" + sp[1] + "\t" + " ".join([str(x) for x in list(feats)]) + '\n')
ffeat.close()