├── README.md ├── calculate_mAP.py ├── data └── readme.txt ├── demo.py └── logs └── useless.txt /README.md: -------------------------------------------------------------------------------- 1 | # DSH_tensorflow 2 | Implemement of DEEP SUPERVISED HASHING FOR FAST IMAGE RETRIEVAL_CVPR2016 3 | 4 | 5 | ## Train: 6 | `python demo.py` 7 | 8 | First you need to modify `demo.py` to ensure `TRAIN=True`,and you can modify `HASHING_BITS` to generate any length hashing bit.After training, in dir `/logs` will save trained model. 9 | 10 | ## Test: 11 | After training, you can change `TRAIN=False`,then run `python.demo.py`,result.txt will be generated. 12 | 13 | ## Calculate mAP: 14 | run `python calculate_mAP.py` to get mAP of this result. 15 | -------------------------------------------------------------------------------- /calculate_mAP.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import time 3 | import os 4 | # read train and test binarayCode 5 | CURRENT_DIR = os.getcwd() 6 | 7 | def getNowTime(): 8 | return '['+time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))+']' 9 | 10 | 11 | def getCode(train_codes,train_groudTruth,test_codes,test_groudTruth): 12 | 13 | line_number = 0 14 | with open(CURRENT_DIR+'/result.txt','r') as f: 15 | for line in f: 16 | temp = line.strip().split('\t') 17 | if line_number < 10000: 18 | test_codes.append([i if i==1 else -1 for i in map(int, list(temp[0]))]) 19 | list2 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] 20 | list2[int(temp[1])] = 1 21 | test_groudTruth.append(list2) # get test ground truth(0-9) 22 | else: 23 | train_codes.append([i if i==1 else -1 for i in map(int, list(temp[0]))]) # change to -1, 1 24 | list2 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0] 25 | list2[int(temp[1])] = 1 26 | train_groudTruth.append(list2) # get test ground truth(0-9) 27 | 28 | line_number += 1 29 | print getNowTime(),'read data finish' 30 | 31 | def getHammingDist(code_a,code_b): 32 | dist = 0 33 | for i in range(len(code_a)): 34 | if code_a[i]!=code_b[i]: 35 | dist += 1 36 | return dist 37 | 38 | 39 | if __name__ =='__main__': 40 | print getNowTime(),'start!' 41 | 42 | train_codes = [] 43 | train_groudTruth =[] 44 | 45 | test_codes = [] 46 | test_groudTruth = [] 47 | # get g.t. and binary code 48 | getCode(train_codes,train_groudTruth,test_codes,test_groudTruth) 49 | train_codes = np.array(train_codes) 50 | train_groudTruth = np.array(train_groudTruth) 51 | test_codes = np.array(test_codes) 52 | test_groudTruth = np.array(test_groudTruth) 53 | numOfTest = 10000 54 | 55 | # generate hanmming martix, g.t. martix 10000*50000 56 | gt_martix = np.dot(test_groudTruth, np.transpose(train_groudTruth)) 57 | print getNowTime(),'gt_martix finish!' 58 | ham_martix = np.dot(test_codes, np.transpose(train_codes)) # hanmming distance map to dot value 59 | print getNowTime(),'ham_martix finish!' 60 | 61 | # sort hanmming martix,Returns the indices that would sort an array. 62 | sorted_ham_martix_index = np.argsort(ham_martix,axis=1) 63 | 64 | # calculate mAP 65 | print getNowTime(),'sort ham_matrix finished,start calculate mAP' 66 | 67 | apall = np.zeros((numOfTest,1),np.float64) 68 | for i in range(numOfTest): 69 | x = 0.0 70 | p = 0 71 | test_oneLine = sorted_ham_martix_index[i,:] 72 | length = test_oneLine.shape[0] 73 | num_return_NN = 5000 # top 1000 74 | for j in range(num_return_NN): 75 | if gt_martix[i][test_oneLine[length-j-1]] == 1: # reverse 76 | x += 1 77 | p += x/(j+1) 78 | if p == 0: 79 | apall[i]=0 80 | else: 81 | apall[i]=p/x 82 | 83 | mAP = np.mean(apall) 84 | print getNowTime(),'mAP:',mAP 85 | 86 | 87 | -------------------------------------------------------------------------------- /data/readme.txt: -------------------------------------------------------------------------------- 1 | download cifar10 python version and unzip it in here. 2 | download url: http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz 3 | -------------------------------------------------------------------------------- /demo.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import tensorflow as tf 4 | import cPickle 5 | 6 | BATCH_SIZE = 200 7 | LR = 0.001 # Learning rate 8 | EPOCH = 600 9 | LOAD_MODEL = False # Whether or not continue train from saved model 10 | TRAIN = True 11 | HASHING_BITS = 12 12 | CURRENT_DIR = os.getcwd() 13 | 14 | 15 | def bias(name, shape, bias_start=0.0, trainable=True): 16 | dtype = tf.float32 17 | var = tf.get_variable(name, shape, tf.float32, trainable=trainable, 18 | initializer=tf.constant_initializer( 19 | bias_start, dtype=dtype)) 20 | return var 21 | 22 | 23 | def weight(name, shape, stddev=0.02, trainable=True): 24 | dtype = tf.float32 25 | var = tf.get_variable(name, shape, tf.float32, trainable=trainable, 26 | initializer=tf.random_normal_initializer( 27 | stddev=stddev, dtype=dtype)) 28 | return var 29 | 30 | 31 | def fully_connected(value, output_shape, name='fully_connected', with_w=False): 32 | value = tf.reshape(value, [BATCH_SIZE, -1]) 33 | shape = value.get_shape().as_list() 34 | 35 | with tf.variable_scope(name): 36 | weights = weight('weights', [shape[1], output_shape], 0.02) 37 | biases = bias('biases', [output_shape], 0.0) 38 | 39 | if with_w: 40 | return tf.matmul(value, weights) + biases, weights, biases 41 | else: 42 | return tf.matmul(value, weights) + biases 43 | 44 | 45 | 46 | def relu(value, name='relu'): 47 | with tf.variable_scope(name): 48 | return tf.nn.relu(value) 49 | 50 | 51 | 52 | def conv2d(value, output_dim, k_h=5, k_w=5, 53 | strides=[1, 1, 1, 1], name='conv2d'): 54 | with tf.variable_scope(name): 55 | weights = weight('weights', 56 | [k_h, k_w, value.get_shape()[-1], output_dim]) 57 | conv = tf.nn.conv2d(value, weights, strides=strides, padding='SAME') 58 | biases = bias('biases', [output_dim]) 59 | conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape()) 60 | 61 | return conv 62 | 63 | 64 | 65 | def pool(value, k_size=[1,3,3,1], 66 | strides=[1, 2, 2, 1], name='pool1'): 67 | with tf.variable_scope(name): 68 | pool = tf.nn.max_pool(value, ksize=k_size, strides=strides, padding='VALID') 69 | return pool 70 | 71 | def pool_avg(value, k_size=[1,3,3,1], 72 | strides=[1, 2, 2, 1], name='pool1'): 73 | with tf.variable_scope(name): 74 | pool = tf.nn.avg_pool(value, ksize=k_size, strides=strides, padding='VALID') 75 | return pool 76 | 77 | def lrn(value, depth_radius=1,alpha=5e-05,beta=0.75, name='lrn1'): 78 | with tf.variable_scope(name): 79 | norm1 = tf.nn.lrn(value, depth_radius=depth_radius, bias=1.0, alpha=alpha, beta=beta) 80 | return norm1 81 | 82 | 83 | 84 | 85 | def discriminator(image, hashing_bits,reuse=False, name='discriminator'): 86 | with tf.name_scope(name): 87 | if reuse: 88 | tf.get_variable_scope().reuse_variables() 89 | conv1 = conv2d(image, output_dim=32, name='d_conv1') 90 | 91 | relu1 = relu(pool(conv1, name='d_pool1'), name='d_relu1') 92 | conv2 = conv2d(lrn(relu1,name='d_lrn1'), output_dim=32, name='d_conv2') 93 | relu2 = relu(pool_avg(conv2, name='d_pool2'), name='d_relu2') 94 | conv3 = conv2d(lrn(relu2, name='d_lrn2'), output_dim=64, name='d_conv3') 95 | pool3 = pool_avg(relu(conv3, name='d_relu3'), name='d_pool3') 96 | relu_ip1 = relu(fully_connected(pool3,output_shape=500, name='d_ip1'), name='d_relu4') 97 | ip2 = fully_connected(relu_ip1,output_shape=hashing_bits, name='d_ip2') 98 | 99 | return ip2 100 | 101 | 102 | 103 | def read_cifar10_data(): 104 | data_dir = CURRENT_DIR+'/data/cifar-10-batches-py/' 105 | train_name = 'data_batch_' 106 | test_name = 'test_batch' 107 | train_X = None 108 | train_Y = None 109 | test_X = None 110 | test_Y = None 111 | 112 | # train data 113 | for i in range(1,6): 114 | file_path = data_dir+train_name+str(i) 115 | with open(file_path, 'rb') as fo: 116 | dict = cPickle.load(fo) 117 | if train_X is None: 118 | train_X = dict['data'] 119 | train_Y = dict['labels'] 120 | else: 121 | train_X = np.concatenate((train_X, dict['data']), axis=0) 122 | train_Y = np.concatenate((train_Y, dict['labels']), axis=0) 123 | # test_data 124 | file_path = data_dir + test_name 125 | with open(file_path, 'rb') as fo: 126 | dict = cPickle.load(fo) 127 | test_X = dict['data'] 128 | test_Y = dict['labels'] 129 | train_X = train_X.reshape((50000, 3, 32, 32)).transpose(0, 2, 3, 1).astype(np.float) 130 | # train_Y = train_Y.reshape((50000)).astype(np.float) 131 | test_X = test_X.reshape((10000, 3, 32, 32)).transpose(0, 2, 3, 1).astype(np.float) 132 | # test_Y.reshape((10000)).astype(np.float) 133 | 134 | train_y_vec = np.zeros((len(train_Y), 10), dtype=np.float) 135 | test_y_vec = np.zeros((len(test_Y), 10), dtype=np.float) 136 | for i, label in enumerate(train_Y): 137 | train_y_vec[i, int(train_Y[i])] = 1. # y_vec[1,3] means #2 row, #4column 138 | for i, label in enumerate(test_Y): 139 | test_y_vec[i, int(test_Y[i])] = 1. # y_vec[1,3] means #2 row, #4column 140 | 141 | return train_X/255., train_y_vec, test_X/255., test_y_vec 142 | 143 | def hashing_loss(image,label,alpha,m): 144 | D = discriminator(image,HASHING_BITS) 145 | w_label = tf.matmul(label,label,False,True) 146 | 147 | r = tf.reshape(tf.reduce_sum(D*D,1),[-1,1]) 148 | p2_distance = r - 2*tf.matmul(D,D,False,True)+tf.transpose(r) 149 | temp = w_label*p2_distance + (1-w_label)*tf.maximum(m-p2_distance,0) 150 | 151 | regularizer = tf.reduce_sum(tf.abs(tf.abs(D) - 1)) 152 | d_loss = tf.reduce_sum(temp)/(BATCH_SIZE*(BATCH_SIZE-1)) + alpha * regularizer/BATCH_SIZE 153 | return d_loss 154 | 155 | 156 | def train(): 157 | train_dir = CURRENT_DIR + '/logs/' 158 | global_step = tf.Variable(0, name='global_step', trainable=False) 159 | image = tf.placeholder(tf.float32, [BATCH_SIZE, 32,32,3], name='image') 160 | label = tf.placeholder(tf.float32, [BATCH_SIZE,10], name='label') 161 | 162 | alpha = tf.constant(0.01,dtype=tf.float32,name='tradeoff') 163 | # set m = 2*HASHING_BITS 164 | m = tf.constant(HASHING_BITS*2,dtype=tf.float32,name='bi_margin') 165 | d_loss_real = hashing_loss(image,label,alpha,m) 166 | 167 | 168 | t_vars = tf.trainable_variables() 169 | d_vars = [var for var in t_vars if 'd_' in var.name] 170 | 171 | saver = tf.train.Saver() 172 | 173 | d_optim = tf.train.AdamOptimizer(LR, beta1=0.5) \ 174 | .minimize(d_loss_real, var_list=d_vars, global_step=global_step) 175 | 176 | os.environ['CUDA_VISIBLE_DEVICES'] = str(0) 177 | config = tf.ConfigProto() 178 | config.gpu_options.per_process_gpu_memory_fraction = 0.4 179 | sess = tf.InteractiveSession(config=config) 180 | 181 | init = tf.global_variables_initializer() 182 | sess.run(init) 183 | 184 | start = 0 185 | if LOAD_MODEL: 186 | print(" [*] Reading checkpoints...") 187 | ckpt = tf.train.get_checkpoint_state(train_dir) 188 | 189 | if ckpt and ckpt.model_checkpoint_path: 190 | ckpt_name = os.path.basename(ckpt.model_checkpoint_path) 191 | saver.restore(sess, os.path.join(train_dir, ckpt_name)) 192 | global_step = ckpt.model_checkpoint_path.split('/')[-1] \ 193 | .split('-')[-1] 194 | print('Loading success, global_step is %s' % global_step) 195 | 196 | start = int(global_step) 197 | 198 | train_x, train_y, test_x, test_y = read_cifar10_data() 199 | 200 | for epoch in range(start,EPOCH): 201 | 202 | batch_idxs = 50000/BATCH_SIZE 203 | 204 | for idx in range(start, batch_idxs): 205 | image_idx = train_x[idx*BATCH_SIZE:(idx+1)*BATCH_SIZE] 206 | label_idx = train_y[idx*BATCH_SIZE:(idx+1)*BATCH_SIZE] 207 | 208 | sess.run([d_optim],feed_dict={image: image_idx, label:label_idx}) 209 | # writer.add_summary(summary_str, idx + 1) 210 | 211 | errD_real = d_loss_real.eval(feed_dict={image: image_idx, label:label_idx}) 212 | 213 | if (idx+1) % 10 == 0: 214 | print("[%3d/%3d][%4d/%4d] d_loss: %.8f" % (epoch+1,EPOCH,idx+1, batch_idxs, errD_real)) 215 | 216 | if (epoch+1) % 10 == 0: 217 | checkpoint_path = os.path.join(train_dir,'my_DSH_model.ckpt') 218 | saver.save(sess, checkpoint_path, global_step=epoch + 1) 219 | print '********* model saved *********' 220 | 221 | 222 | sess.close() 223 | 224 | def toBinaryString(binary_like_values): 225 | numOfImage,bit_length = binary_like_values.shape 226 | list_string_binary = [] 227 | for i in range(numOfImage): 228 | str = '' 229 | for j in range(bit_length): 230 | str += '0' if binary_like_values[i][j] <= 0 else '1' 231 | list_string_binary.append(str) 232 | return list_string_binary 233 | 234 | def evaluate(): 235 | checkpoint_dir = CURRENT_DIR + '/logs/' 236 | 237 | image = tf.placeholder(tf.float32, [BATCH_SIZE, 32, 32, 3], name='image') 238 | D = discriminator(image,HASHING_BITS) 239 | res = tf.sign(D) 240 | 241 | print("Reading checkpoints...") 242 | ckpt = tf.train.get_checkpoint_state(checkpoint_dir) 243 | saver = tf.train.Saver(tf.all_variables()) 244 | os.environ['CUDA_VISIBLE_DEVICES'] = str(0) 245 | config = tf.ConfigProto() 246 | config.gpu_options.per_process_gpu_memory_fraction = 0.8 247 | sess = tf.InteractiveSession(config=config) 248 | 249 | train_x, train_y, test_x, test_y = read_cifar10_data() 250 | file_res = open('result.txt','w') 251 | # sys.stdout = file_res 252 | if ckpt and ckpt.model_checkpoint_path: 253 | ckpt_name = os.path.basename(ckpt.model_checkpoint_path) 254 | global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1] 255 | saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name)) 256 | print('Loading success, global_step is %s' % global_step) 257 | for i in range(10000/BATCH_SIZE): 258 | eval_sess = sess.run(D, feed_dict={image: test_x[i*BATCH_SIZE:(i+1)*BATCH_SIZE]}) 259 | # print(eval_sess) 260 | w_res = toBinaryString(eval_sess) 261 | w_label = np.argmax(test_y[i*BATCH_SIZE:(i+1)*BATCH_SIZE],axis=1) 262 | for j in range(BATCH_SIZE): 263 | file_res.write(w_res[j]+'\t'+str(w_label[j])+'\n') 264 | for i in range(50000/BATCH_SIZE): 265 | eval_sess = sess.run(D, feed_dict={image: train_x[i*BATCH_SIZE:(i+1)*BATCH_SIZE]}) 266 | # print(eval_sess) 267 | w_res = toBinaryString(eval_sess) 268 | w_label = np.argmax(train_y[i*BATCH_SIZE:(i+1)*BATCH_SIZE],axis=1) 269 | for j in range(BATCH_SIZE): 270 | file_res.write(w_res[j]+'\t'+str(w_label[j])+'\n') 271 | 272 | 273 | # eval_sess = sess.run(res, feed_dict={image: test_x[:BATCH_SIZE]}) 274 | # eval_sess = sess.run(res, feed_dict={image: test_x[:BATCH_SIZE]}) 275 | # print(eval_sess) 276 | file_res.close() 277 | sess.close() 278 | 279 | if __name__ == '__main__': 280 | 281 | if TRAIN: 282 | train() 283 | else: 284 | evaluate() 285 | -------------------------------------------------------------------------------- /logs/useless.txt: -------------------------------------------------------------------------------- 1 | just want to add dir logs/ 2 | --------------------------------------------------------------------------------