├── model └── readme ├── DCGAN-tensorflow ├── DCGAN.png ├── average.png ├── dcgan.yml ├── resizeImage.py ├── LICENSE ├── README.md ├── ops.py ├── main.py ├── download.py ├── utils.py └── model.py ├── changeIndex.py ├── random_erasing.py ├── evaluate_rerank.py ├── prepare.py ├── evaluate.py ├── model.py ├── re_ranking.py ├── README.md ├── test.py └── train_baseline.py /model/readme: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/DCGAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qiaoguan/Person-reid-GAN-pytorch/HEAD/DCGAN-tensorflow/DCGAN.png -------------------------------------------------------------------------------- /DCGAN-tensorflow/average.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/qiaoguan/Person-reid-GAN-pytorch/HEAD/DCGAN-tensorflow/average.png -------------------------------------------------------------------------------- /DCGAN-tensorflow/dcgan.yml: -------------------------------------------------------------------------------- 1 | name: dcgan 2 | channels: 3 | - https://conda.anaconda.org/menpo 4 | - conda-forge 5 | dependencies: 6 | - python==3.6.1 7 | - numpy 8 | - matplotlib 9 | - pillow 10 | - h5py 11 | - tqdm 12 | - imageio 13 | - scipy 14 | - protobuf 15 | - pip: 16 | - moviepy 17 | - tensorflow-gpu==0.12.1 18 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/resizeImage.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | 4 | if __name__=='__main__': 5 | fileout='resize_market' 6 | filein='gen_market' 7 | if not os.path.isdir(fileout): 8 | os.mkdir(fileout) 9 | for name in os.listdir(filein): 10 | img=Image.open(filein+'/'+name) 11 | out=img.resize((64,128),Image.ANTIALIAS) 12 | out.save(fileout+'/'+name) -------------------------------------------------------------------------------- /changeIndex.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # change name of the folder(e.g. 0002,0007,0010,0011... to 0,1,2,3) 3 | import numpy as np 4 | import os 5 | from shutil import copyfile 6 | 7 | original_path='/home/gq123/guanqiao/deeplearning/reid/market/pytorch' 8 | 9 | #copy folder tree from source to destination 10 | def copyfolder(src,dst): 11 | files=os.listdir(src) 12 | if not os.path.isdir(dst): 13 | os.mkdir(dst) 14 | for tt in files: 15 | copyfile(src+'/'+tt,dst+'/'+tt) 16 | 17 | train_save_path = original_path + '/val_new' 18 | data_path=original_path+'/val' 19 | if not os.path.isdir(train_save_path): 20 | os.mkdir(train_save_path) 21 | 22 | reid_index=0 23 | folders=os.listdir(data_path) 24 | for foldernames in folders: 25 | copyfolder(data_path+'/'+foldernames,train_save_path+'/'+str(reid_index).zfill(4)) 26 | reid_index=reid_index+1 27 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Taehoon Kim 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/README.md: -------------------------------------------------------------------------------- 1 | # DCGAN in Tensorflow 2 | 3 | ## Prerequisites 4 | 5 | - Python 2.7 or Python 3.3+ 6 | - [Tensorflow 0.12.1](https://github.com/tensorflow/tensorflow/tree/r0.12) (Notice that it is not the latest version) 7 | - [SciPy](http://www.scipy.org/install.html) 8 | - [pillow](https://github.com/python-pillow/Pillow) 9 | - CUDA 8.0 10 | 11 | Add Cuda Path to bashrc first 12 | ```bash 13 | export LD_LIBRARY_PATH="/usr/local/cuda-8.0/lib64:$LD_LIBRARY_PATH" 14 | ``` 15 | 16 | We recommend you to install anaconda. Here we write a simple script for you to install the dependence by anaconda. 17 | ```python 18 | # install env (especially for old version Tensorflow) 19 | conda env create -f dcgan.yml 20 | # activate env, then you can run code in this env without downgrading the outside Tensorflow. 21 | source activate dcgan 22 | ``` 23 | 24 | ### Let's start 25 | 26 | ### 1.Train 27 | ```bash 28 | mkdir data 29 | copy your dataset(market) to data folder 30 | python main.py --dataset market --options 1 31 | ``` 32 | `market` is the dir path which contains images. You can change it to your dataset path. 33 | 34 | ### 2.Test 35 | ```bash 36 | python main.py --dataset market --options 5 --output_path gen_market --sample_size 24000 37 | python resizeImage.py 38 | ``` 39 | It will use your trained model and generate 24000 images for the following semi-supervised training. the generated images are stored in gen_market folder, then run resizeImage.py, the generated images will be resized into 126*64 stored in gen_0000 market. after that, add gen_0000 to the training set of market, so the generated images will be used as training image to help model training. 40 | -------------------------------------------------------------------------------- /random_erasing.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from torchvision.transforms import * 4 | 5 | from PIL import Image 6 | import random 7 | import math 8 | import numpy as np 9 | import torch 10 | 11 | class RandomErasing(object): 12 | """ Randomly selects a rectangle region in an image and erases its pixels. 13 | 'Random Erasing Data Augmentation' by Zhong et al. 14 | See https://arxiv.org/pdf/1708.04896.pdf 15 | Args: 16 | probability: The probability that the Random Erasing operation will be performed. 17 | sl: Minimum proportion of erased area against input image. 18 | sh: Maximum proportion of erased area against input image. 19 | r1: Minimum aspect ratio of erased area. 20 | mean: Erasing value. 21 | """ 22 | 23 | def __init__(self, probability = 0.5, sl = 0.02, sh = 0.4, r1 = 0.3, mean=[0.4914, 0.4822, 0.4465]): 24 | self.probability = probability 25 | self.mean = mean 26 | self.sl = sl 27 | self.sh = sh 28 | self.r1 = r1 29 | 30 | def __call__(self, img): 31 | 32 | if random.uniform(0, 1) > self.probability: 33 | return img 34 | 35 | for attempt in range(100): 36 | area = img.size()[1] * img.size()[2] 37 | 38 | target_area = random.uniform(self.sl, self.sh) * area 39 | aspect_ratio = random.uniform(self.r1, 1/self.r1) 40 | 41 | h = int(round(math.sqrt(target_area * aspect_ratio))) 42 | w = int(round(math.sqrt(target_area / aspect_ratio))) 43 | 44 | if w < img.size()[2] and h < img.size()[1]: 45 | x1 = random.randint(0, img.size()[1] - h) 46 | y1 = random.randint(0, img.size()[2] - w) 47 | if img.size()[0] == 3: 48 | img[0, x1:x1+h, y1:y1+w] = self.mean[0] 49 | img[1, x1:x1+h, y1:y1+w] = self.mean[1] 50 | img[2, x1:x1+h, y1:y1+w] = self.mean[2] 51 | else: 52 | img[0, x1:x1+h, y1:y1+w] = self.mean[0] 53 | return img 54 | 55 | return img 56 | -------------------------------------------------------------------------------- /evaluate_rerank.py: -------------------------------------------------------------------------------- 1 | import scipy.io 2 | import torch 3 | import numpy as np 4 | import time 5 | from re_ranking import re_ranking 6 | ####################################################################### 7 | # Evaluate 8 | def evaluate(score,ql,qc,gl,gc): 9 | index = np.argsort(score) #from small to large 10 | #index = index[::-1] 11 | # good index 12 | query_index = np.argwhere(gl==ql) 13 | camera_index = np.argwhere(gc==qc) 14 | 15 | good_index = np.setdiff1d(query_index, camera_index, assume_unique=True) 16 | junk_index1 = np.argwhere(gl==-1) 17 | junk_index2 = np.intersect1d(query_index, camera_index) 18 | junk_index = np.append(junk_index2, junk_index1) #.flatten()) 19 | 20 | CMC_tmp = compute_mAP(index, good_index, junk_index) 21 | return CMC_tmp 22 | 23 | 24 | def compute_mAP(index, good_index, junk_index): 25 | ap = 0 26 | cmc = torch.IntTensor(len(index)).zero_() 27 | if good_index.size==0: # if empty 28 | cmc[0] = -1 29 | return ap,cmc 30 | 31 | # remove junk_index 32 | mask = np.in1d(index, junk_index, invert=True) 33 | index = index[mask] 34 | 35 | # find good_index index 36 | ngood = len(good_index) 37 | mask = np.in1d(index, good_index) 38 | rows_good = np.argwhere(mask==True) 39 | rows_good = rows_good.flatten() 40 | 41 | cmc[rows_good[0]:] = 1 42 | for i in range(ngood): 43 | d_recall = 1.0/ngood 44 | precision = (i+1)*1.0/(rows_good[i]+1) 45 | if rows_good[i]!=0: 46 | old_precision = i*1.0/rows_good[i] 47 | else: 48 | old_precision=1.0 49 | ap = ap + d_recall*(old_precision + precision)/2 50 | 51 | return ap, cmc 52 | 53 | ###################################################################### 54 | result = scipy.io.loadmat('pytorch_result.mat') 55 | query_feature = result['query_f'] 56 | query_cam = result['query_cam'][0] 57 | query_label = result['query_label'][0] 58 | gallery_feature = result['gallery_f'] 59 | gallery_cam = result['gallery_cam'][0] 60 | gallery_label = result['gallery_label'][0] 61 | 62 | CMC = torch.IntTensor(len(gallery_label)).zero_() 63 | ap = 0.0 64 | #re-ranking 65 | print('calculate initial distance') 66 | q_g_dist = np.dot(query_feature, np.transpose(gallery_feature)) 67 | q_q_dist = np.dot(query_feature, np.transpose(query_feature)) 68 | g_g_dist = np.dot(gallery_feature, np.transpose(gallery_feature)) 69 | since = time.time() 70 | re_rank = re_ranking(q_g_dist, q_q_dist, g_g_dist) 71 | time_elapsed = time.time() - since 72 | print('Reranking complete in {:.0f}m {:.0f}s'.format( 73 | time_elapsed // 60, time_elapsed % 60)) 74 | for i in range(len(query_label)): 75 | ap_tmp, CMC_tmp = evaluate(re_rank[i,:],query_label[i],query_cam[i],gallery_label,gallery_cam) 76 | if CMC_tmp[0]==-1: 77 | continue 78 | CMC = CMC + CMC_tmp 79 | ap += ap_tmp 80 | #print(i, CMC_tmp[0]) 81 | 82 | CMC = CMC.float() 83 | CMC = CMC/len(query_label) #average CMC 84 | print('top1:%f top5:%f top10:%f mAP:%f'%(CMC[0],CMC[4],CMC[9],ap/len(query_label))) 85 | -------------------------------------------------------------------------------- /prepare.py: -------------------------------------------------------------------------------- 1 | import os 2 | from shutil import copyfile 3 | 4 | # You only need to change this line to your dataset download path 5 | download_path = '/home/gq123/guanqiao/deeplearning/reid_baseline_pytorch/market' 6 | 7 | if not os.path.isdir(download_path): 8 | print('please change the download_path') 9 | 10 | save_path = download_path + '/pytorch' 11 | if not os.path.isdir(save_path): 12 | os.mkdir(save_path) 13 | #----------------------------------------- 14 | #query 15 | query_path = download_path + '/query' 16 | query_save_path = download_path + '/pytorch/query' 17 | if not os.path.isdir(query_save_path): 18 | os.mkdir(query_save_path) 19 | 20 | for root, dirs, files in os.walk(query_path, topdown=True): 21 | for name in files: 22 | if not name[-3:]=='jpg': 23 | continue 24 | ID = name.split('_') 25 | src_path = query_path + '/' + name 26 | dst_path = query_save_path + '/' + ID[0] 27 | if not os.path.isdir(dst_path): 28 | os.mkdir(dst_path) 29 | copyfile(src_path, dst_path + '/' + name) 30 | 31 | #----------------------------------------- 32 | #gallery 33 | gallery_path = download_path + '/bounding_box_test' 34 | gallery_save_path = download_path + '/pytorch/gallery' 35 | if not os.path.isdir(gallery_save_path): 36 | os.mkdir(gallery_save_path) 37 | 38 | for root, dirs, files in os.walk(gallery_path, topdown=True): 39 | for name in files: 40 | if not name[-3:]=='jpg': 41 | continue 42 | ID = name.split('_') 43 | src_path = gallery_path + '/' + name 44 | dst_path = gallery_save_path + '/' + ID[0] 45 | if not os.path.isdir(dst_path): 46 | os.mkdir(dst_path) 47 | copyfile(src_path, dst_path + '/' + name) 48 | 49 | #--------------------------------------- 50 | #train_all 51 | train_path = download_path + '/bounding_box_train' 52 | train_save_path = download_path + '/pytorch/train_all' 53 | if not os.path.isdir(train_save_path): 54 | os.mkdir(train_save_path) 55 | 56 | for root, dirs, files in os.walk(train_path, topdown=True): 57 | for name in files: 58 | if not name[-3:]=='jpg': 59 | continue 60 | ID = name.split('_') 61 | src_path = train_path + '/' + name 62 | dst_path = train_save_path + '/' + ID[0] 63 | if not os.path.isdir(dst_path): 64 | os.mkdir(dst_path) 65 | copyfile(src_path, dst_path + '/' + name) 66 | 67 | 68 | #--------------------------------------- 69 | #train_val 70 | train_path = download_path + '/bounding_box_train' 71 | train_save_path = download_path + '/pytorch/train' 72 | val_save_path = download_path + '/pytorch/val' 73 | if not os.path.isdir(train_save_path): 74 | os.mkdir(train_save_path) 75 | os.mkdir(val_save_path) 76 | 77 | for root, dirs, files in os.walk(train_path, topdown=True): 78 | for name in files: 79 | if not name[-3:]=='jpg': 80 | continue 81 | ID = name.split('_') 82 | src_path = train_path + '/' + name 83 | dst_path = train_save_path + '/' + ID[0] 84 | if not os.path.isdir(dst_path): 85 | os.mkdir(dst_path) 86 | dst_path = val_save_path + '/' + ID[0] #first image is used as val image, the other images for train 87 | os.mkdir(dst_path) 88 | copyfile(src_path, dst_path + '/' + name) 89 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import scipy.io 2 | import torch 3 | import numpy as np 4 | import time 5 | 6 | ####################################################################### 7 | # Evaluate find the corresponding person in the gallery set 8 | # qf: 1*1024, ql:1, qc:1, gf: 19732*1024, gl: 19732, gc: 19732 9 | def evaluate(qf,ql,qc,gf,gl,gc): 10 | query = qf 11 | score = np.dot(gf,query) # cosine distance 12 | # predict index 13 | index = np.argsort(score) #from small to large(small equals similar) 14 | index = index[::-1] # reverse 19732*1 15 | #index = index[0:2000] 16 | # good index 17 | query_index = np.argwhere(gl==ql) # 59*1 18 | camera_index = np.argwhere(gc==qc) # 3156*1 19 | 20 | good_index = np.setdiff1d(query_index, camera_index, assume_unique=True) # the difference of the set 51*1 21 | junk_index1 = np.argwhere(gl==-1) # 3819*1 22 | junk_index2 = np.intersect1d(query_index, camera_index) # find the intersection of two arrays # 8*1 23 | junk_index = np.append(junk_index2, junk_index1) #.flatten()) # 3827*1 24 | 25 | CMC_tmp = compute_mAP(index, good_index, junk_index) 26 | return CMC_tmp 27 | 28 | 29 | def compute_mAP(index, good_index, junk_index): 30 | ap = 0 31 | cmc = torch.IntTensor(len(index)).zero_() # create a zero tensor 32 | if good_index.size==0: # if empty 33 | cmc[0] = -1 34 | return ap,cmc 35 | 36 | # remove junk_index 37 | mask = np.in1d(index, junk_index, invert=True) 38 | index = index[mask] # remove junk_index from index 39 | 40 | # find good_index index 41 | ngood = len(good_index) 42 | mask = np.in1d(index, good_index) 43 | rows_good = np.argwhere(mask==True) # np.argwhere, find the index of element which satisfy the condition that mask is True 44 | rows_good = rows_good.flatten() 45 | ''' 46 | print(index.shape) 47 | print(ngood) 48 | print(rows_good.shape) 49 | name=raw_input() 50 | ''' 51 | cmc[rows_good[0]:] = 1 52 | for i in range(ngood): 53 | d_recall = 1.0/ngood 54 | precision = (i+1)*1.0/(rows_good[i]+1) 55 | if rows_good[i]!=0: 56 | old_precision = i*1.0/rows_good[i] 57 | else: 58 | old_precision=1.0 59 | ap = ap + d_recall*(old_precision + precision)/2 60 | 61 | return ap, cmc 62 | 63 | ###################################################################### 64 | result = scipy.io.loadmat('pytorch_result.mat') 65 | query_feature = result['query_f'] # 3368*1024 66 | query_cam = result['query_cam'][0] # 3368 67 | query_label = result['query_label'][0] # 3368 68 | gallery_feature = result['gallery_f'] #19732*1024 69 | gallery_cam = result['gallery_cam'][0] #19732 70 | gallery_label = result['gallery_label'][0] # 19732 71 | 72 | CMC = torch.IntTensor(len(gallery_label)).zero_() 73 | ap = 0.0 74 | #print(query_label) len=3368 75 | ''' 76 | print('---------------------') 77 | print(len(query_feature)) 78 | print(len(query_label)) 79 | print(len(query_cam)) 80 | print(len(gallery_feature)) 81 | print(len(gallery_label)) 82 | print(len(gallery_cam)) 83 | ''' 84 | for i in range(len(query_label)): 85 | ap_tmp, CMC_tmp = evaluate(query_feature[i],query_label[i],query_cam[i],gallery_feature,gallery_label,gallery_cam) 86 | if CMC_tmp[0]==-1: 87 | continue 88 | CMC = CMC + CMC_tmp 89 | ap += ap_tmp 90 | # print(i, CMC_tmp[0]) 91 | 92 | CMC = CMC.float() 93 | CMC = CMC/len(query_label) #average CMC 94 | print('top1:%f top5:%f top10:%f mAP:%f'%(CMC[0],CMC[4],CMC[9],ap/len(query_label))) -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | from torchvision import models 5 | from torch.autograd import Variable 6 | 7 | ###################################################################### 8 | def weights_init_kaiming(m): 9 | classname = m.__class__.__name__ 10 | # print(classname) 11 | if classname.find('Conv') != -1: 12 | init.kaiming_normal(m.weight.data, a=0, mode='fan_in') 13 | elif classname.find('Linear') != -1: 14 | init.kaiming_normal(m.weight.data, a=0, mode='fan_out') 15 | init.constant(m.bias.data, 0.0) 16 | elif classname.find('BatchNorm1d') != -1: 17 | init.normal(m.weight.data, 1.0, 0.02) 18 | init.constant(m.bias.data, 0.0) 19 | 20 | def weights_init_classifier(m): 21 | classname = m.__class__.__name__ 22 | if classname.find('Linear') != -1: 23 | init.normal(m.weight.data, std=0.001) 24 | init.constant(m.bias.data, 0.0) 25 | 26 | class ft_net(nn.Module): 27 | 28 | def __init__(self, class_num ): 29 | super(ft_net,self).__init__() 30 | model_ft = models.resnet50(pretrained=True) 31 | 32 | # avg pooling to global pooling 33 | model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1)) 34 | 35 | num_ftrs = model_ft.fc.in_features # extract feature parameters of fully collected layers 36 | add_block = [] 37 | num_bottleneck = 512 38 | add_block += [nn.Linear(num_ftrs, num_bottleneck)] # add a linear layer, batchnorm layer, leakyrelu layer and dropout layer 39 | add_block += [nn.BatchNorm1d(num_bottleneck)] 40 | add_block += [nn.LeakyReLU(0.1)] 41 | add_block += [nn.Dropout(p=0.5)] #default dropout rate 0.5 42 | #transforms.CenterCrop(224), 43 | add_block = nn.Sequential(*add_block) 44 | add_block.apply(weights_init_kaiming) 45 | model_ft.fc = add_block 46 | self.model = model_ft 47 | 48 | classifier = [] 49 | classifier += [nn.Linear(num_bottleneck, class_num)] # class_num classification 50 | classifier = nn.Sequential(*classifier) 51 | classifier.apply(weights_init_classifier) 52 | self.classifier = classifier 53 | 54 | def forward(self, x): 55 | x = self.model(x) 56 | x = self.classifier(x) 57 | return x 58 | 59 | 60 | class ft_net_dense(nn.Module): 61 | def __init__(self, class_num ): 62 | super(ft_net_dense,self).__init__() 63 | model_ft = models.densenet121(pretrained=True) 64 | # add pooling to the model 65 | # in the originial version, pooling is written in the forward function 66 | model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1,1)) 67 | 68 | add_block = [] 69 | num_bottleneck = 512 70 | add_block += [nn.Linear(1024, num_bottleneck)] #For ResNet, it is 2048 71 | add_block += [nn.BatchNorm1d(num_bottleneck)] 72 | add_block += [nn.LeakyReLU(0.1)] 73 | add_block += [nn.Dropout(p=0.5)] 74 | add_block = nn.Sequential(*add_block) 75 | add_block.apply(weights_init_kaiming) 76 | model_ft.fc = add_block 77 | self.model = model_ft 78 | 79 | classifier = [] 80 | classifier += [nn.Linear(num_bottleneck, class_num)] 81 | classifier = nn.Sequential(*classifier) 82 | classifier.apply(weights_init_classifier) 83 | self.classifier = classifier 84 | 85 | def forward(self, x): 86 | x = self.model.features(x) 87 | x = x.view(x.size(0),-1) 88 | x = self.model.fc(x) 89 | x = self.classifier(x) 90 | return x 91 | 92 | # debug model structure 93 | #net = ft_net(751) 94 | 95 | #net = ft_net_dense(751) 96 | #print(net) 97 | ''' 98 | input = Variable(torch.FloatTensor(8, 3, 224, 224)) 99 | output = net(input) 100 | print('net output size:') 101 | #print(output.shape) 102 | ''' -------------------------------------------------------------------------------- /DCGAN-tensorflow/ops.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | import tensorflow as tf 4 | 5 | from tensorflow.python.framework import ops 6 | 7 | from utils import * 8 | 9 | try: 10 | image_summary = tf.image_summary 11 | scalar_summary = tf.scalar_summary 12 | histogram_summary = tf.histogram_summary 13 | merge_summary = tf.merge_summary 14 | SummaryWriter = tf.train.SummaryWriter 15 | except: 16 | image_summary = tf.summary.image 17 | scalar_summary = tf.summary.scalar 18 | histogram_summary = tf.summary.histogram 19 | merge_summary = tf.summary.merge 20 | SummaryWriter = tf.summary.FileWriter 21 | 22 | if "concat_v2" in dir(tf): 23 | def concat(tensors, axis, *args, **kwargs): 24 | return tf.concat_v2(tensors, axis, *args, **kwargs) 25 | else: 26 | def concat(tensors, axis, *args, **kwargs): 27 | return tf.concat(tensors, axis, *args, **kwargs) 28 | 29 | class batch_norm(object): 30 | def __init__(self, epsilon=1e-5, momentum = 0.9, name="batch_norm"): 31 | with tf.variable_scope(name): 32 | self.epsilon = epsilon 33 | self.momentum = momentum 34 | self.name = name 35 | 36 | def __call__(self, x, train=True): 37 | return tf.contrib.layers.batch_norm(x, 38 | decay=self.momentum, 39 | updates_collections=None, 40 | epsilon=self.epsilon, 41 | scale=True, 42 | is_training=train, 43 | scope=self.name) 44 | 45 | def conv_cond_concat(x, y): 46 | """Concatenate conditioning vector on feature map axis.""" 47 | x_shapes = x.get_shape() 48 | y_shapes = y.get_shape() 49 | return concat([ 50 | x, y*tf.ones([x_shapes[0], x_shapes[1], x_shapes[2], y_shapes[3]])], 3) 51 | 52 | def conv2d(input_, output_dim, 53 | k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02, 54 | name="conv2d"): 55 | with tf.variable_scope(name): 56 | w = tf.get_variable('w', [k_h, k_w, input_.get_shape()[-1], output_dim], 57 | initializer=tf.truncated_normal_initializer(stddev=stddev)) 58 | conv = tf.nn.conv2d(input_, w, strides=[1, d_h, d_w, 1], padding='SAME') 59 | 60 | biases = tf.get_variable('biases', [output_dim], initializer=tf.constant_initializer(0.0)) 61 | conv = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape()) 62 | 63 | return conv 64 | 65 | def deconv2d(input_, output_shape, 66 | k_h=5, k_w=5, d_h=2, d_w=2, stddev=0.02, 67 | name="deconv2d", with_w=False): 68 | with tf.variable_scope(name): 69 | # filter : [height, width, output_channels, in_channels] 70 | w = tf.get_variable('w', [k_h, k_w, output_shape[-1], input_.get_shape()[-1]], 71 | initializer=tf.random_normal_initializer(stddev=stddev)) 72 | 73 | try: 74 | deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, 75 | strides=[1, d_h, d_w, 1]) 76 | 77 | # Support for verisons of TensorFlow before 0.7.0 78 | except AttributeError: 79 | deconv = tf.nn.deconv2d(input_, w, output_shape=output_shape, 80 | strides=[1, d_h, d_w, 1]) 81 | 82 | biases = tf.get_variable('biases', [output_shape[-1]], initializer=tf.constant_initializer(0.0)) 83 | deconv = tf.reshape(tf.nn.bias_add(deconv, biases), deconv.get_shape()) 84 | 85 | if with_w: 86 | return deconv, w, biases 87 | else: 88 | return deconv 89 | 90 | def lrelu(x, leak=0.2, name="lrelu"): 91 | return tf.maximum(x, leak*x) 92 | 93 | def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, with_w=False): 94 | shape = input_.get_shape().as_list() 95 | 96 | with tf.variable_scope(scope or "Linear"): 97 | matrix = tf.get_variable("Matrix", [shape[1], output_size], tf.float32, 98 | tf.random_normal_initializer(stddev=stddev)) 99 | bias = tf.get_variable("bias", [output_size], 100 | initializer=tf.constant_initializer(bias_start)) 101 | if with_w: 102 | return tf.matmul(input_, matrix) + bias, matrix, bias 103 | else: 104 | return tf.matmul(input_, matrix) + bias 105 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import scipy.misc 3 | import numpy as np 4 | 5 | from model import DCGAN 6 | from utils import pp, visualize, to_json, show_all_variables 7 | 8 | import tensorflow as tf 9 | 10 | flags = tf.app.flags 11 | flags.DEFINE_integer("epoch", 25, "Epoch to train [25]") 12 | flags.DEFINE_float("learning_rate", 0.0002, "Learning rate of for adam [0.0002]") 13 | flags.DEFINE_float("beta1", 0.5, "Momentum term of adam [0.5]") 14 | flags.DEFINE_integer("train_size", np.inf, "The size of train images [np.inf]") 15 | flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]") 16 | flags.DEFINE_integer("input_height", 128, "The size of image to use (will be center cropped). [108]") 17 | flags.DEFINE_integer("input_width", None, "The size of image to use (will be center cropped). If None, same value as input_height [None]") 18 | flags.DEFINE_integer("output_height", 128, "The size of the output images to produce [64]") 19 | flags.DEFINE_integer("output_width", None, "The size of the output images to produce. If None, same value as output_height [None]") 20 | flags.DEFINE_string("output_path", 'duke_result',"output image path") 21 | flags.DEFINE_integer("sample_size", 1000,"How much sample you want to output") 22 | flags.DEFINE_integer("options", 1,"output option") 23 | flags.DEFINE_string("dataset", "celebA", "The name of dataset [celebA, mnist, lsun]") 24 | flags.DEFINE_string("input_fname_pattern", "*.jpg", "Glob pattern of filename of input images [*]") 25 | flags.DEFINE_string("checkpoint_dir", "checkpoint", "Directory name to save the checkpoints [checkpoint]") 26 | flags.DEFINE_string("sample_dir", "samples", "Directory name to save the image samples [samples]") 27 | flags.DEFINE_boolean("train", False, "True for training, False for testing [False]") 28 | flags.DEFINE_boolean("crop", False, "True for training, False for testing [False]") 29 | flags.DEFINE_boolean("visualize", False, "True for visualizing, False for nothing [False]") 30 | FLAGS = flags.FLAGS 31 | 32 | def main(_): 33 | pp.pprint(flags.FLAGS.__flags) 34 | 35 | if FLAGS.input_width is None: 36 | FLAGS.input_width = FLAGS.input_height 37 | if FLAGS.output_width is None: 38 | FLAGS.output_width = FLAGS.output_height 39 | 40 | if not os.path.exists(FLAGS.checkpoint_dir): 41 | os.makedirs(FLAGS.checkpoint_dir) 42 | if not os.path.exists(FLAGS.sample_dir): 43 | os.makedirs(FLAGS.sample_dir) 44 | 45 | #gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333) 46 | run_config = tf.ConfigProto() 47 | run_config.gpu_options.allow_growth=True 48 | 49 | with tf.Session(config=run_config) as sess: 50 | if FLAGS.dataset == 'mnist': 51 | dcgan = DCGAN( 52 | sess, 53 | input_width=FLAGS.input_width, 54 | input_height=FLAGS.input_height, 55 | output_width=FLAGS.output_width, 56 | output_height=FLAGS.output_height, 57 | batch_size=FLAGS.batch_size, 58 | sample_num=FLAGS.batch_size, 59 | y_dim=10, 60 | dataset_name=FLAGS.dataset, 61 | input_fname_pattern=FLAGS.input_fname_pattern, 62 | crop=FLAGS.crop, 63 | checkpoint_dir=FLAGS.checkpoint_dir, 64 | sample_dir=FLAGS.sample_dir) 65 | else: 66 | dcgan = DCGAN( 67 | sess, 68 | input_width=FLAGS.input_width, 69 | input_height=FLAGS.input_height, 70 | output_width=FLAGS.output_width, 71 | output_height=FLAGS.output_height, 72 | batch_size=FLAGS.batch_size, 73 | sample_num=FLAGS.batch_size, 74 | dataset_name=FLAGS.dataset, 75 | input_fname_pattern=FLAGS.input_fname_pattern, 76 | crop=FLAGS.crop, 77 | checkpoint_dir=FLAGS.checkpoint_dir, 78 | sample_dir=FLAGS.sample_dir) 79 | 80 | show_all_variables() 81 | 82 | if FLAGS.train: 83 | dcgan.train(FLAGS) 84 | else: 85 | if not dcgan.load(FLAGS.checkpoint_dir)[0]: 86 | raise Exception("[!] Train a model first, then run test mode") 87 | 88 | 89 | # to_json("./web/js/layers.js", [dcgan.h0_w, dcgan.h0_b, dcgan.g_bn0], 90 | # [dcgan.h1_w, dcgan.h1_b, dcgan.g_bn1], 91 | # [dcgan.h2_w, dcgan.h2_b, dcgan.g_bn2], 92 | # [dcgan.h3_w, dcgan.h3_b, dcgan.g_bn3], 93 | # [dcgan.h4_w, dcgan.h4_b, None]) 94 | 95 | # Below is codes for visualization 96 | OPTION = FLAGS.options 97 | visualize(sess, dcgan, FLAGS, OPTION) 98 | 99 | if __name__ == '__main__': 100 | tf.app.run() 101 | -------------------------------------------------------------------------------- /re_ranking.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2/python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Mon Jun 26 14:46:56 2017 5 | @author: luohao 6 | Modified by Houjing Huang, 2017-12-22. 7 | - This version accepts distance matrix instead of raw features. 8 | - The difference of `/` division between python 2 and 3 is handled. 9 | - numpy.float16 is replaced by numpy.float32 for numerical precision. 10 | 11 | Modified by Zhedong Zheng, 2018-1-12. 12 | - replace sort with topK, which save about 30s. 13 | """ 14 | 15 | """ 16 | CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. 17 | url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf 18 | Matlab version: https://github.com/zhunzhong07/person-re-ranking 19 | """ 20 | 21 | """ 22 | API 23 | q_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery] 24 | q_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query] 25 | g_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery] 26 | k1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3) 27 | Returns: 28 | final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery] 29 | """ 30 | 31 | 32 | import numpy as np 33 | 34 | def k_reciprocal_neigh( initial_rank, i, k1): 35 | forward_k_neigh_index = initial_rank[i,:k1+1] 36 | backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1] 37 | fi = np.where(backward_k_neigh_index==i)[0] 38 | return forward_k_neigh_index[fi] 39 | 40 | def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3): 41 | # The following naming, e.g. gallery_num, is different from outer scope. 42 | # Don't care about it. 43 | original_dist = np.concatenate( 44 | [np.concatenate([q_q_dist, q_g_dist], axis=1), 45 | np.concatenate([q_g_dist.T, g_g_dist], axis=1)], 46 | axis=0) 47 | original_dist = 2. - 2 * original_dist #np.power(original_dist, 2).astype(np.float32) 48 | original_dist = np.transpose(1. * original_dist/np.max(original_dist,axis = 0)) 49 | V = np.zeros_like(original_dist).astype(np.float32) 50 | #initial_rank = np.argsort(original_dist).astype(np.int32) 51 | # top K1+1 52 | initial_rank = np.argpartition( original_dist, range(1,k1+1) ) 53 | 54 | query_num = q_g_dist.shape[0] 55 | all_num = original_dist.shape[0] 56 | 57 | for i in range(all_num): 58 | # k-reciprocal neighbors 59 | k_reciprocal_index = k_reciprocal_neigh( initial_rank, i, k1) 60 | k_reciprocal_expansion_index = k_reciprocal_index 61 | for j in range(len(k_reciprocal_index)): 62 | candidate = k_reciprocal_index[j] 63 | candidate_k_reciprocal_index = k_reciprocal_neigh( initial_rank, candidate, int(np.around(k1/2))) 64 | if len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index))> 2./3*len(candidate_k_reciprocal_index): 65 | k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index) 66 | 67 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) 68 | weight = np.exp(-original_dist[i,k_reciprocal_expansion_index]) 69 | V[i,k_reciprocal_expansion_index] = 1.*weight/np.sum(weight) 70 | 71 | original_dist = original_dist[:query_num,] 72 | if k2 != 1: 73 | V_qe = np.zeros_like(V,dtype=np.float32) 74 | for i in range(all_num): 75 | V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:],axis=0) 76 | V = V_qe 77 | del V_qe 78 | del initial_rank 79 | invIndex = [] 80 | for i in range(all_num): 81 | invIndex.append(np.where(V[:,i] != 0)[0]) 82 | 83 | jaccard_dist = np.zeros_like(original_dist,dtype = np.float32) 84 | 85 | for i in range(query_num): 86 | temp_min = np.zeros(shape=[1,all_num],dtype=np.float32) 87 | indNonZero = np.where(V[i,:] != 0)[0] 88 | indImages = [] 89 | indImages = [invIndex[ind] for ind in indNonZero] 90 | for j in range(len(indNonZero)): 91 | temp_min[0,indImages[j]] = temp_min[0,indImages[j]]+ np.minimum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]]) 92 | jaccard_dist[i] = 1-temp_min/(2.-temp_min) 93 | 94 | final_dist = jaccard_dist*(1-lambda_value) + original_dist*lambda_value 95 | del original_dist 96 | del V 97 | del jaccard_dist 98 | final_dist = final_dist[:query_num,query_num:] 99 | return final_dist 100 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Person-reid-GAN-pytorch 2 | A Pytorch Implementation of ["Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro"(ICCV17)](http://openaccess.thecvf.com/content_ICCV_2017/papers/Zheng_Unlabeled_Samples_Generated_ICCV_2017_paper.pdf), the official code is available [here](https://github.com/layumi/Person-reID_GAN)(in matlab). 3 | 4 | 5 | We arrived **Rank@1=93.55%, mAP=90.67%** only with a very easy model. 6 | 7 | 8 | **Random Erasing** is added to help train as a data augmentation method. the details of Random Erasing is available [here](https://github.com/zhunzhong07/Random-Erasing) 9 | 10 | **re-rank strategy** is used to deal with the initial result, the details of re-rank method is available [here](https://github.com/zhunzhong07/person-re-ranking) 11 | 12 | 13 | ## Model Structure（we simply alter the model from ResNet and DenseNet） 14 | You may learn more from `model.py`. 15 | We add one linear layer(bottleneck), one batchnorm layer and relu. 16 | 17 | ## Prerequisites 18 | 19 | - Python 2.7 20 | - GPU 21 | - Numpy 22 | - Pytorch 23 | - Torchvision 24 | 25 | ## Getting started 26 | ### Installation 27 | - Install Pytorch(the version is 0.2.0_3) from http://pytorch.org/ 28 | - Install Torchvision from the source 29 | ``` 30 | git clone https://github.com/pytorch/vision 31 | cd vision 32 | python setup.py install 33 | ``` 34 | ## Dataset & Preparation 35 | Download [Market1501 Dataset](http://www.liangzheng.org/Project/project_reid.html) 36 | 37 | Preparation: Put the images with the same id in one folder. You may use 38 | ```bash 39 | python prepare.py 40 | ``` 41 | Remember to change the dataset path to your own path. 42 | 43 | Preparation: change the name of the folder to 0 to n-1(where n is the number of class, i.e. the number of person), the folder name is the label of each person (all the pictures under the same folder are the same person) 44 | ```bash 45 | python changeIndex.py 46 | ``` 47 | the usage of the DCGAN model is under the DCGAN folder, first use market1501 to train the dcgan model, then generated some pictures using the trained DCGAN model, then you can use different numbers of generated images to help train the model using LSRO. 48 | 49 | **the generated images are in the gen_0000 folder, you can copy this folder under your training set, for more details, you can refer to DCGAN-TENSORFLOW folder** 50 | 51 | Our baseline code is only finetuned from resNet or DenseNet, 52 | we use pretained DenseNet as baseline to train our model, the archieved result are as follows: 53 | 54 | |Rank@1 | mAP | Note| 55 | | ----- | ---- | ---- | 56 | |0.921 | 0.793| ----| 57 | |0.934   |  0.907|re-rank| 58 | 59 | using LSRO loss and added some image generated by DCGAN model, the achieved result are as follow: 60 | 61 | |**Batchsize|Multi/Single GPU training |Rank@1 | mAP | Note**| 62 | | ----- | ---- | ----- | ---- | ---- | 63 | | 32 | Single | 0.9162 | 0.7887 | add 0 generated image| 64 | | 32 | Single | 0.9355 | 0.9067 | after re-rank| 65 | | 32 | Multi | 0.8367 | 0.6442 | add 0 generated image| 66 | | 32 | Multi | 0.8655 | 0.8143 | after re-rank| 67 | | 64 | Multi | 0.843 | 0.646 | add 0 generated image| 68 | | 64 | Multi | 0.872 | 0.815 | after re-rank| 69 | | 32 | Single | 0.919 | 0.798 | add 6000 generated image| 70 | | 32 | Single | 0.932 | 0.9012 | after re-rank| 71 | | 64 | Multi | 0.909 | 0.779 | add 12000 generated image| 72 | | 64 | Multi | 0.931 | 0.896 | after re-rank| 73 | | 32 | Single | 0.925 | 0.801 | add 12000 generated image| 74 | | 32 | Single | 0.939 | 0.904 | after re-rank| 75 | | 64 | Multi | 0.915 | 0.790 | add 18000 generated image| 76 | | 64 | Multi | 0.933 | 0.899 | after re-rank| 77 | | 64 | Multi | 0.909 | 0.773 | add 24000 generated image| 78 | | 64 | Multi | 0.924 | 0.887 | after re-rank| 79 | | 32 | Single | 0.918 | 0.790 | add 24000 generated image| 80 | | 32 | Single | 0.932 | 0.899 | after re-rank| 81 | 82 | To save trained model, we make a dir. 83 | ```bash 84 | mkdir model 85 | ``` 86 | 87 | 88 | ## Train 89 | Train the baseline by 90 | ```bash 91 | python train_baseline.py --use_dense 92 | ``` 93 | 94 | `--name` the name of model.(ResNet or DesNet) 95 | 96 | `--data_dir` the path of the training data. 97 | 98 | `--batchsize` batch size. 99 | 100 | `--erasing_p` random erasing probability. 101 | 102 | 103 | ## Test 104 | Use trained model to extract feature by 105 | ```bash 106 | python test.py --which_epoch 99 --use_dense 107 | ``` 108 | `--gpu_ids` which gpu to run. 109 | 110 | `--name` the dir name of trained model. 111 | 112 | `--which_epoch` select the i-th model. 113 | 114 | `--data_dir` the path of the testing data. 115 | 116 | `--batchsize` batch size. 117 | 118 | 119 | ## Evaluation 120 | ```bash 121 | python evaluate.py 122 | ``` 123 | It will output Rank@1, Rank@5, Rank@10 and mAP results. 124 | 125 | For mAP calculation, you also can refer to the [C++ code for Oxford Building](http://www.robots.ox.ac.uk/~vgg/data/oxbuildings/compute_ap.cpp). We use the triangle mAP calculation (consistent with the Market1501 original code). 126 | 127 | ### re-ranking 128 | ```bash 129 | python evaluate_rerank.py 130 | ``` 131 | **It may take more than 10G Memory to run.** So run it on a powerful machine if possible. 132 | 133 | It will output Rank@1, Rank@5, Rank@10 and mAP results. 134 | 135 | ### Conclusion 136 | 137 | **when the baseline result is not so high, the generated images can help model training(see multi-gpu training, add GAN images VS not add) , thus can improve the performance(more robust) , while when the baseline result is high(rank-1, 0.934), its difficult to improve the result. when batchsize is set to 32, the result is the best, and single-gpu training achieves better result than multi-gpu training.** 138 | 139 | ## Thanks 140 | 141 | * Many, many thanks to [layumi](https://github.com/layumi/Person-reID_GAN) for his Great work! 142 | 143 | ### ~~p.s. If you have any questions, you can open an issue!~~ This Repo will no longer be supported for personal reasons! 144 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/download.py: -------------------------------------------------------------------------------- 1 | """ 2 | Modification of https://github.com/stanfordnlp/treelstm/blob/master/scripts/download.py 3 | 4 | Downloads the following: 5 | - Celeb-A dataset 6 | - LSUN dataset 7 | - MNIST dataset 8 | """ 9 | 10 | from __future__ import print_function 11 | import os 12 | import sys 13 | import gzip 14 | import json 15 | import shutil 16 | import zipfile 17 | import argparse 18 | import requests 19 | import subprocess 20 | from tqdm import tqdm 21 | from six.moves import urllib 22 | 23 | parser = argparse.ArgumentParser(description='Download dataset for DCGAN.') 24 | parser.add_argument('datasets', metavar='N', type=str, nargs='+', choices=['celebA', 'lsun', 'mnist'], 25 | help='name of dataset to download [celebA, lsun, mnist]') 26 | 27 | def download(url, dirpath): 28 | filename = url.split('/')[-1] 29 | filepath = os.path.join(dirpath, filename) 30 | u = urllib.request.urlopen(url) 31 | f = open(filepath, 'wb') 32 | filesize = int(u.headers["Content-Length"]) 33 | print("Downloading: %s Bytes: %s" % (filename, filesize)) 34 | 35 | downloaded = 0 36 | block_sz = 8192 37 | status_width = 70 38 | while True: 39 | buf = u.read(block_sz) 40 | if not buf: 41 | print('') 42 | break 43 | else: 44 | print('', end='\r') 45 | downloaded += len(buf) 46 | f.write(buf) 47 | status = (("[%-" + str(status_width + 1) + "s] %3.2f%%") % 48 | ('=' * int(float(downloaded) / filesize * status_width) + '>', downloaded * 100. / filesize)) 49 | print(status, end='') 50 | sys.stdout.flush() 51 | f.close() 52 | return filepath 53 | 54 | def download_file_from_google_drive(id, destination): 55 | URL = "https://docs.google.com/uc?export=download" 56 | session = requests.Session() 57 | 58 | response = session.get(URL, params={ 'id': id }, stream=True) 59 | token = get_confirm_token(response) 60 | 61 | if token: 62 | params = { 'id' : id, 'confirm' : token } 63 | response = session.get(URL, params=params, stream=True) 64 | 65 | save_response_content(response, destination) 66 | 67 | def get_confirm_token(response): 68 | for key, value in response.cookies.items(): 69 | if key.startswith('download_warning'): 70 | return value 71 | return None 72 | 73 | def save_response_content(response, destination, chunk_size=32*1024): 74 | total_size = int(response.headers.get('content-length', 0)) 75 | with open(destination, "wb") as f: 76 | for chunk in tqdm(response.iter_content(chunk_size), total=total_size, 77 | unit='B', unit_scale=True, desc=destination): 78 | if chunk: # filter out keep-alive new chunks 79 | f.write(chunk) 80 | 81 | def unzip(filepath): 82 | print("Extracting: " + filepath) 83 | dirpath = os.path.dirname(filepath) 84 | with zipfile.ZipFile(filepath) as zf: 85 | zf.extractall(dirpath) 86 | os.remove(filepath) 87 | 88 | def download_celeb_a(dirpath): 89 | data_dir = 'celebA' 90 | if os.path.exists(os.path.join(dirpath, data_dir)): 91 | print('Found Celeb-A - skip') 92 | return 93 | 94 | filename, drive_id = "img_align_celeba.zip", "0B7EVK8r0v71pZjFTYXZWM3FlRnM" 95 | save_path = os.path.join(dirpath, filename) 96 | 97 | if os.path.exists(save_path): 98 | print('[*] {} already exists'.format(save_path)) 99 | else: 100 | download_file_from_google_drive(drive_id, save_path) 101 | 102 | zip_dir = '' 103 | with zipfile.ZipFile(save_path) as zf: 104 | zip_dir = zf.namelist()[0] 105 | zf.extractall(dirpath) 106 | os.remove(save_path) 107 | os.rename(os.path.join(dirpath, zip_dir), os.path.join(dirpath, data_dir)) 108 | 109 | def _list_categories(tag): 110 | url = 'http://lsun.cs.princeton.edu/htbin/list.cgi?tag=' + tag 111 | f = urllib.request.urlopen(url) 112 | return json.loads(f.read()) 113 | 114 | def _download_lsun(out_dir, category, set_name, tag): 115 | url = 'http://lsun.cs.princeton.edu/htbin/download.cgi?tag={tag}' \ 116 | '&category={category}&set={set_name}'.format(**locals()) 117 | print(url) 118 | if set_name == 'test': 119 | out_name = 'test_lmdb.zip' 120 | else: 121 | out_name = '{category}_{set_name}_lmdb.zip'.format(**locals()) 122 | out_path = os.path.join(out_dir, out_name) 123 | cmd = ['curl', url, '-o', out_path] 124 | print('Downloading', category, set_name, 'set') 125 | subprocess.call(cmd) 126 | 127 | def download_lsun(dirpath): 128 | data_dir = os.path.join(dirpath, 'lsun') 129 | if os.path.exists(data_dir): 130 | print('Found LSUN - skip') 131 | return 132 | else: 133 | os.mkdir(data_dir) 134 | 135 | tag = 'latest' 136 | #categories = _list_categories(tag) 137 | categories = ['bedroom'] 138 | 139 | for category in categories: 140 | _download_lsun(data_dir, category, 'train', tag) 141 | _download_lsun(data_dir, category, 'val', tag) 142 | _download_lsun(data_dir, '', 'test', tag) 143 | 144 | def download_mnist(dirpath): 145 | data_dir = os.path.join(dirpath, 'mnist') 146 | if os.path.exists(data_dir): 147 | print('Found MNIST - skip') 148 | return 149 | else: 150 | os.mkdir(data_dir) 151 | url_base = 'http://yann.lecun.com/exdb/mnist/' 152 | file_names = ['train-images-idx3-ubyte.gz', 153 | 'train-labels-idx1-ubyte.gz', 154 | 't10k-images-idx3-ubyte.gz', 155 | 't10k-labels-idx1-ubyte.gz'] 156 | for file_name in file_names: 157 | url = (url_base+file_name).format(**locals()) 158 | print(url) 159 | out_path = os.path.join(data_dir,file_name) 160 | cmd = ['curl', url, '-o', out_path] 161 | print('Downloading ', file_name) 162 | subprocess.call(cmd) 163 | cmd = ['gzip', '-d', out_path] 164 | print('Decompressing ', file_name) 165 | subprocess.call(cmd) 166 | 167 | def prepare_data_dir(path = './data'): 168 | if not os.path.exists(path): 169 | os.mkdir(path) 170 | 171 | if __name__ == '__main__': 172 | args = parser.parse_args() 173 | prepare_data_dir() 174 | 175 | if any(name in args.datasets for name in ['CelebA', 'celebA', 'celebA']): 176 | download_celeb_a('./data') 177 | if 'lsun' in args.datasets: 178 | download_lsun('./data') 179 | if 'mnist' in args.datasets: 180 | download_mnist('./data') 181 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | ''' 3 | if the model is trained by multi-GPU, use the upper load_network() function, else use the load_network() below. 4 | ''' 5 | from __future__ import print_function, division 6 | 7 | import argparse 8 | import torch 9 | import torch.nn as nn 10 | import torch.optim as optim 11 | from torch.optim import lr_scheduler 12 | from torch.autograd import Variable 13 | import numpy as np 14 | import torchvision 15 | from torchvision import datasets, models, transforms 16 | import time 17 | import os 18 | import scipy.io 19 | from model import ft_net, ft_net_dense 20 | 21 | ###################################################################### 22 | # Options 23 | # -------- 24 | parser = argparse.ArgumentParser(description='Training') 25 | parser.add_argument('--gpu_ids',default='3', type=str,help='gpu_ids: e.g. 0 0,1,2 0,2') 26 | parser.add_argument('--which_epoch',default='best', type=str, help='0,1,2,3...or last') 27 | parser.add_argument('--test_dir',default='/home/gq123/guanqiao/deeplearning/reid/market/pytorch',type=str, help='./test_data') 28 | parser.add_argument('--name', default='ft_DesNet121', type=str, help='save model path') 29 | parser.add_argument('--batchsize', default=32, type=int, help='batchsize') 30 | parser.add_argument('--use_dense', action='store_true', help='use densenet121' ) 31 | 32 | opt = parser.parse_args() 33 | 34 | str_ids = opt.gpu_ids.split(',') 35 | #which_epoch = opt.which_epoch 36 | name = opt.name 37 | test_dir = opt.test_dir 38 | 39 | gpu_ids = [] 40 | for str_id in str_ids: 41 | id = int(str_id) 42 | if id >=0: 43 | gpu_ids.append(id) 44 | 45 | # set gpu ids 46 | if len(gpu_ids)>0: 47 | torch.cuda.set_device(gpu_ids[0]) 48 | 49 | ###################################################################### 50 | # Load Data 51 | # --------- 52 | # 53 | # We will use torchvision and torch.utils.data packages for loading the 54 | # data. 55 | # 56 | data_transforms = transforms.Compose([ 57 | transforms.Resize((288,144), interpolation=3), 58 | transforms.ToTensor(), 59 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 60 | ############### Ten Crop 61 | #transforms.TenCrop(224), 62 | #transforms.Lambda(lambda crops: torch.stack( 63 | # [transforms.ToTensor()(crop) 64 | # for crop in crops] 65 | # )), 66 | #transforms.Lambda(lambda crops: torch.stack( 67 | # [transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(crop) 68 | # for crop in crops] 69 | # )) 70 | ]) 71 | 72 | 73 | data_dir = test_dir 74 | image_datasets = {x: datasets.ImageFolder( os.path.join(data_dir,x) ,data_transforms) for x in ['gallery','query']} 75 | dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=opt.batchsize, 76 | shuffle=False, num_workers=4) for x in ['gallery','query']} 77 | 78 | class_names = image_datasets['query'].classes 79 | use_gpu = torch.cuda.is_available() 80 | 81 | ###################################################################### 82 | # Load model 83 | ''' 84 | #----------single gpu training----------------- 85 | def load_network(network): 86 | save_path = os.path.join('./model',name,'net_%s.pth'%opt.which_epoch) 87 | network.load_state_dict(torch.load(save_path)) 88 | return network 89 | ''' 90 | 91 | #-----multi-gpu training--------- 92 | def load_network(network): 93 | save_path = os.path.join('./model',name,'net_%s.pth'%opt.which_epoch) 94 | state_dict = torch.load(save_path) 95 | # create new OrderedDict that does not contain `module.` 96 | from collections import OrderedDict 97 | new_state_dict = OrderedDict() 98 | for k, v in state_dict.items(): 99 | namekey = k[7:] # remove `module.` 100 | new_state_dict[namekey] = v 101 | # load params 102 | network.load_state_dict(new_state_dict) 103 | return network 104 | 105 | ###################################################################### 106 | # Extract feature 107 | # ---------------------- 108 | # 109 | # Extract feature from a trained model. 110 | # 111 | def fliplr(img): 112 | '''flip horizontal''' 113 | inv_idx = torch.arange(img.size(3)-1,-1,-1).long() # N x C x H x W 114 | img_flip = img.index_select(3,inv_idx) 115 | return img_flip 116 | 117 | def extract_feature(model,dataloaders): 118 | features = torch.FloatTensor() 119 | count = 0 120 | for data in dataloaders: 121 | img, label = data 122 | n, c, h, w = img.size() 123 | count += n 124 | # print(count) 125 | if opt.use_dense: 126 | ff = torch.FloatTensor(n,1024).zero_() 127 | else: 128 | ff = torch.FloatTensor(n,2048).zero_() 129 | for i in range(2): 130 | if(i==1): 131 | img = fliplr(img) 132 | input_img = Variable(img.cuda()) 133 | outputs = model(input_img) 134 | f = outputs.data.cpu() 135 | #print(f.size()) 136 | ff = ff+f 137 | # norm feature 138 | fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) # L2 normalize 139 | ff = ff.div(fnorm.expand_as(ff)) 140 | features = torch.cat((features,ff), 0) 141 | return features 142 | 143 | def get_id(img_path): 144 | camera_id = [] 145 | labels = [] 146 | for path, v in img_path: 147 | filename = path.split('/')[-1] 148 | label = filename[0:4] 149 | camera = filename.split('c')[1] 150 | if label[0:2]=='-1': 151 | labels.append(-1) 152 | else: 153 | labels.append(int(label)) 154 | camera_id.append(int(camera[0])) 155 | return camera_id, labels 156 | 157 | gallery_path = image_datasets['gallery'].imgs 158 | query_path = image_datasets['query'].imgs 159 | 160 | gallery_cam,gallery_label = get_id(gallery_path) 161 | query_cam,query_label = get_id(query_path) 162 | 163 | ###################################################################### 164 | # Load Collected data Trained model 165 | print('-------test-----------') 166 | if opt.use_dense: 167 | model_structure = ft_net_dense(751) 168 | else: 169 | model_structure = ft_net(751) 170 | model = load_network(model_structure) 171 | 172 | # Remove the final fc layer and classifier layer 173 | model.model.fc = nn.Sequential() 174 | model.classifier = nn.Sequential() 175 | 176 | # Change to test mode 177 | model = model.eval() 178 | if use_gpu: 179 | model = model.cuda() 180 | 181 | # Extract feature 182 | gallery_feature = extract_feature(model,dataloaders['gallery']) 183 | query_feature = extract_feature(model,dataloaders['query']) 184 | 185 | # Save to Matlab for check 186 | result = {'gallery_f':gallery_feature.numpy(),'gallery_label':gallery_label,'gallery_cam':gallery_cam,'query_f':query_feature.numpy(),'query_label':query_label,'query_cam':query_cam} 187 | scipy.io.savemat('pytorch_result.mat',result) 188 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Some codes from https://github.com/Newmu/dcgan_code 3 | """ 4 | from __future__ import division 5 | import math 6 | import json 7 | import os 8 | import random 9 | import pprint 10 | import scipy.misc 11 | import numpy as np 12 | from time import gmtime, strftime 13 | from six.moves import xrange 14 | 15 | import tensorflow as tf 16 | import tensorflow.contrib.slim as slim 17 | 18 | pp = pprint.PrettyPrinter() 19 | 20 | get_stddev = lambda x, k_h, k_w: 1/math.sqrt(k_w*k_h*x.get_shape()[-1]) 21 | 22 | def show_all_variables(): 23 | model_vars = tf.trainable_variables() 24 | slim.model_analyzer.analyze_vars(model_vars, print_info=True) 25 | 26 | def get_image(image_path, input_height, input_width, 27 | resize_height=64, resize_width=64, 28 | crop=True, grayscale=False): 29 | image = imread(image_path, grayscale) 30 | return transform(image, input_height, input_width, 31 | resize_height, resize_width, crop) 32 | 33 | def save_images(images, size, image_path): 34 | return imsave(inverse_transform(images), size, image_path) 35 | 36 | def save_images2(images, size, image_path): 37 | return scipy.misc.imsave(image_path, scipy.misc.imresize(images, size, 'bilinear')) 38 | 39 | def imread(path, grayscale = False): 40 | if (grayscale): 41 | return scipy.misc.imread(path, flatten = True).astype(np.float) 42 | else: 43 | return scipy.misc.imread(path).astype(np.float) 44 | 45 | def merge_images(images, size): 46 | return inverse_transform(images) 47 | 48 | def merge(images, size): 49 | h, w = images.shape[1], images.shape[2] 50 | if (images.shape[3] in (3,4)): 51 | c = images.shape[3] 52 | img = np.zeros((h * size[0], w * size[1], c)) 53 | for idx, image in enumerate(images): 54 | i = idx % size[1] 55 | j = idx // size[1] 56 | img[j * h:j * h + h, i * w:i * w + w, :] = image 57 | return img 58 | elif images.shape[3]==1: 59 | img = np.zeros((h * size[0], w * size[1])) 60 | for idx, image in enumerate(images): 61 | i = idx % size[1] 62 | j = idx // size[1] 63 | img[j * h:j * h + h, i * w:i * w + w] = image[:,:,0] 64 | return img 65 | else: 66 | raise ValueError('in merge(images,size) images parameter ' 67 | 'must have dimensions: HxW or HxWx3 or HxWx4') 68 | 69 | def imsave(images, size, path): 70 | image = np.squeeze(merge(images, size)) 71 | return scipy.misc.imsave(path, image) 72 | 73 | def center_crop(x, crop_h, crop_w, 74 | resize_h=64, resize_w=64): 75 | if crop_w is None: 76 | crop_w = crop_h 77 | h, w = x.shape[:2] 78 | j = int(round((h - crop_h)/2.)) 79 | i = int(round((w - crop_w)/2.)) 80 | return scipy.misc.imresize( 81 | x[j:j+crop_h, i:i+crop_w], [resize_h, resize_w]) 82 | 83 | def transform(image, input_height, input_width, 84 | resize_height=64, resize_width=64, crop=True): 85 | if crop: 86 | cropped_image = center_crop( 87 | image, input_height, input_width, 88 | resize_height, resize_width) 89 | else: 90 | cropped_image = scipy.misc.imresize(image, [resize_height, resize_width]) 91 | return np.array(cropped_image)/127.5 - 1. 92 | 93 | def inverse_transform(images): 94 | return (images+1.)/2. 95 | 96 | def to_json(output_path, *layers): 97 | with open(output_path, "w") as layer_f: 98 | lines = "" 99 | for w, b, bn in layers: 100 | layer_idx = w.name.split('/')[0].split('h')[1] 101 | 102 | B = b.eval() 103 | 104 | if "lin/" in w.name: 105 | W = w.eval() 106 | depth = W.shape[1] 107 | else: 108 | W = np.rollaxis(w.eval(), 2, 0) 109 | depth = W.shape[0] 110 | 111 | biases = {"sy": 1, "sx": 1, "depth": depth, "w": ['%.2f' % elem for elem in list(B)]} 112 | if bn != None: 113 | gamma = bn.gamma.eval() 114 | beta = bn.beta.eval() 115 | 116 | gamma = {"sy": 1, "sx": 1, "depth": depth, "w": ['%.2f' % elem for elem in list(gamma)]} 117 | beta = {"sy": 1, "sx": 1, "depth": depth, "w": ['%.2f' % elem for elem in list(beta)]} 118 | else: 119 | gamma = {"sy": 1, "sx": 1, "depth": 0, "w": []} 120 | beta = {"sy": 1, "sx": 1, "depth": 0, "w": []} 121 | 122 | if "lin/" in w.name: 123 | fs = [] 124 | for w in W.T: 125 | fs.append({"sy": 1, "sx": 1, "depth": W.shape[0], "w": ['%.2f' % elem for elem in list(w)]}) 126 | 127 | lines += """ 128 | var layer_%s = { 129 | "layer_type": "fc", 130 | "sy": 1, "sx": 1, 131 | "out_sx": 1, "out_sy": 1, 132 | "stride": 1, "pad": 0, 133 | "out_depth": %s, "in_depth": %s, 134 | "biases": %s, 135 | "gamma": %s, 136 | "beta": %s, 137 | "filters": %s 138 | };""" % (layer_idx.split('_')[0], W.shape[1], W.shape[0], biases, gamma, beta, fs) 139 | else: 140 | fs = [] 141 | for w_ in W: 142 | fs.append({"sy": 5, "sx": 5, "depth": W.shape[3], "w": ['%.2f' % elem for elem in list(w_.flatten())]}) 143 | 144 | lines += """ 145 | var layer_%s = { 146 | "layer_type": "deconv", 147 | "sy": 5, "sx": 5, 148 | "out_sx": %s, "out_sy": %s, 149 | "stride": 2, "pad": 1, 150 | "out_depth": %s, "in_depth": %s, 151 | "biases": %s, 152 | "gamma": %s, 153 | "beta": %s, 154 | "filters": %s 155 | };""" % (layer_idx, 2**(int(layer_idx)+2), 2**(int(layer_idx)+2), 156 | W.shape[0], W.shape[3], biases, gamma, beta, fs) 157 | layer_f.write(" ".join(lines.replace("'","").split())) 158 | 159 | def make_gif(images, fname, duration=2, true_image=False): 160 | import moviepy.editor as mpy 161 | 162 | def make_frame(t): 163 | try: 164 | x = images[int(len(images)/duration*t)] 165 | except: 166 | x = images[-1] 167 | 168 | if true_image: 169 | return x.astype(np.uint8) 170 | else: 171 | return ((x+1)/2*255).astype(np.uint8) 172 | 173 | clip = mpy.VideoClip(make_frame, duration=duration) 174 | clip.write_gif(fname, fps = len(images) / duration) 175 | 176 | def visualize(sess, dcgan, config, option): 177 | image_frame_dim = int(math.ceil(config.batch_size**.5)) 178 | if option == 0: 179 | z_sample = np.random.uniform(-0.5, 0.5, size=(config.batch_size, dcgan.z_dim)) 180 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}) 181 | save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_%s.png' % strftime("%Y%m%d%H%M%S", gmtime())) 182 | elif option == 1: 183 | values = np.arange(0, 1, 1./config.batch_size) 184 | for idx in xrange(100): 185 | print(" [*] %d" % idx) 186 | z_sample = np.zeros([config.batch_size, dcgan.z_dim]) 187 | for kdx, z in enumerate(z_sample): 188 | z[idx] = values[kdx] 189 | 190 | if config.dataset == "mnist": 191 | y = np.random.choice(10, config.batch_size) 192 | y_one_hot = np.zeros((config.batch_size, 10)) 193 | y_one_hot[np.arange(config.batch_size), y] = 1 194 | 195 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample, dcgan.y: y_one_hot}) 196 | else: 197 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}) 198 | 199 | save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_arange_%s.png' % (idx)) 200 | elif option == 2: 201 | values = np.arange(0, 1, 1./config.batch_size) 202 | for idx in [random.randint(0, 99) for _ in xrange(100)]: 203 | print(" [*] %d" % idx) 204 | z = np.random.uniform(-0.2, 0.2, size=(dcgan.z_dim)) 205 | z_sample = np.tile(z, (config.batch_size, 1)) 206 | #z_sample = np.zeros([config.batch_size, dcgan.z_dim]) 207 | for kdx, z in enumerate(z_sample): 208 | z[idx] = values[kdx] 209 | 210 | if config.dataset == "mnist": 211 | y = np.random.choice(10, config.batch_size) 212 | y_one_hot = np.zeros((config.batch_size, 10)) 213 | y_one_hot[np.arange(config.batch_size), y] = 1 214 | 215 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample, dcgan.y: y_one_hot}) 216 | else: 217 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}) 218 | 219 | try: 220 | make_gif(samples, './samples/test_gif_%s.gif' % (idx)) 221 | except: 222 | save_images(samples, [image_frame_dim, image_frame_dim], './samples/test_%s.png' % strftime("%Y%m%d%H%M%S", gmtime())) 223 | elif option == 3: 224 | values = np.arange(0, 1, 1./config.batch_size) 225 | for idx in xrange(100): 226 | print(" [*] %d" % idx) 227 | z_sample = np.zeros([config.batch_size, dcgan.z_dim]) 228 | for kdx, z in enumerate(z_sample): 229 | z[idx] = values[kdx] 230 | 231 | samples = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}) 232 | make_gif(samples, './samples/test_gif_%s.gif' % (idx)) 233 | elif option == 4: 234 | image_set = [] 235 | values = np.arange(0, 1, 1./config.batch_size) 236 | 237 | for idx in xrange(100): 238 | print(" [*] %d" % idx) 239 | z_sample = np.zeros([config.batch_size, dcgan.z_dim]) 240 | for kdx, z in enumerate(z_sample): z[idx] = values[kdx] 241 | 242 | image_set.append(sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample})) 243 | make_gif(image_set[-1], './samples/test_gif_%s.gif' % (idx)) 244 | 245 | new_image_set = [merge(np.array([images[idx] for images in image_set]), [10, 10]) \ 246 | for idx in range(64) + range(63, -1, -1)] 247 | make_gif(new_image_set, './samples/test_gif_merged.gif', duration=8) 248 | elif option == 5: 249 | for j in xrange(int(config.sample_size/config.batch_size)): 250 | z_sample = np.random.uniform(-1.0,1.0,size=(config.batch_size,dcgan.z_dim)) 251 | image = sess.run(dcgan.sampler, feed_dict={dcgan.z: z_sample}) 252 | if not os.path.isdir(config.output_path): 253 | os.mkdir(config.output_path) 254 | for i in xrange(config.batch_size): 255 | print(" [*] %s" % str(64*j+i)) 256 | save_images2(image[i],[256,256,3,3,3],"./%s/test_%d_%d.jpg"%(config.output_path,j,i)) 257 | 258 | def image_manifold_size(num_images): 259 | manifold_h = int(np.floor(np.sqrt(num_images))) 260 | manifold_w = int(np.ceil(np.sqrt(num_images))) 261 | assert manifold_h * manifold_w == num_images 262 | return manifold_h, manifold_w 263 | -------------------------------------------------------------------------------- /train_baseline.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | #Author: qiaoguan(https://github.com/qiaoguan/Person_reID_baseline_pytorch) 3 | ''' 4 | this is the baseline, if do not add gen_0000 folder(generateed images by DCGAN) under the training set, 5 | so the LSRO equals to crossentropy loss, and the generated_image_size is 0. else the loss function will use the generated images, the loss function for 6 | the generated images and original images are not the same. 7 | ''' 8 | from __future__ import print_function, division 9 | import cv2 10 | import argparse 11 | import torch 12 | import torch.nn as nn 13 | import torch.optim as optim 14 | from torch.optim import lr_scheduler 15 | from torch.autograd import Variable 16 | import numpy as np 17 | import torchvision 18 | from torchvision import datasets, models, transforms 19 | from torchvision.datasets.folder import default_loader 20 | import matplotlib 21 | matplotlib.use('agg') 22 | import matplotlib.pyplot as plt 23 | from PIL import Image 24 | import time 25 | import os 26 | from model import ft_net, ft_net_dense 27 | from random_erasing import RandomErasing 28 | import json 29 | import torch.nn.functional as F 30 | from torch.utils.data import Dataset,DataLoader 31 | 32 | ###################################################################### 33 | # Options 34 | parser = argparse.ArgumentParser(description='Training') 35 | #parser.add_argument('--gpu_ids',default='3', type=str,help='gpu_ids: e.g. 0 0,1,2 0,2') 36 | parser.add_argument('--name',default='ft_DesNet121', type=str, help='output model name') 37 | parser.add_argument('--data_dir',default='/home/gq123/guanqiao/deeplearning/reid/market/pytorch',type=str, help='training dir path') 38 | parser.add_argument('--batchsize', default=64, type=int, help='batchsize') 39 | parser.add_argument('--erasing_p', default=0.8, type=float, help='Random Erasing probability, in [0,1]') 40 | parser.add_argument('--use_dense', action='store_true', help='use densenet121' ) 41 | opt = parser.parse_args() 42 | 43 | data_dir = opt.data_dir 44 | name = opt.name 45 | 46 | generated_image_size=24000 47 | ''' 48 | str_ids = opt.gpu_ids.split(',') 49 | gpu_ids = [] 50 | for str_id in str_ids: 51 | gid = int(str_id) 52 | if gid >=0: 53 | gpu_ids.append(gid) 54 | # set gpu ids 55 | if len(gpu_ids)>0: 56 | torch.cuda.set_device(gpu_ids[0]) 57 | ''' 58 | ###################################################################### 59 | transform_train_list = [ 60 | #transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) 61 | transforms.Resize(144, interpolation=3), 62 | transforms.RandomCrop((256,128)), 63 | # transforms.Resize(256,interpolation=3), 64 | # transforms.RandomCrop(224,224), 65 | transforms.RandomHorizontalFlip(), 66 | transforms.ToTensor(), 67 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 68 | ] 69 | 70 | if opt.erasing_p>0: 71 | transform_train_list = transform_train_list + [RandomErasing(opt.erasing_p)] 72 | 73 | #print(transform_train_list) 74 | 75 | transform_val_list = [ 76 | transforms.Resize(size=(256,128),interpolation=3), #Image.BICUBIC 77 | # transforms.Resize(256,interpolation=3), 78 | # transforms.RandomCrop(224,224), 79 | transforms.ToTensor(), 80 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 81 | ] 82 | 83 | 84 | data_transforms = { 85 | 'train': transforms.Compose( transform_train_list ), 86 | 'val': transforms.Compose(transform_val_list), 87 | } 88 | 89 | # read dcgan data 90 | class dcganDataset(Dataset): 91 | def __init__(self, root,transform=None, targte_transform=None): 92 | super(dcganDataset,self).__init__() 93 | self.image_dir = os.path.join(opt.data_dir, root) 94 | self.samples=[] # train_data xxx_label_flag_yyy.jpg 95 | self.img_label=[] 96 | self.img_flag=[] 97 | self.transform=transform 98 | self.targte_transform=targte_transform 99 | # self.class_num=len(os.listdir(self.image_dir)) # the number of the class 100 | self.train_val=root # judge whether it is used for training for testing 101 | if root=='train_new' : 102 | for folder in os.listdir(self.image_dir): 103 | fdir=self.image_dir+'/'+folder # folder gen_0000 means the images are generated images, so their flags are 1 104 | if folder == 'gen_0000': 105 | for files in os.listdir(fdir): 106 | temp=folder+'_'+files 107 | self.img_label.append(int(folder[-4:])) 108 | self.img_flag.append(1) 109 | self.samples.append(temp) 110 | else: 111 | for files in os.listdir(fdir): 112 | temp=folder+'_'+files 113 | self.img_label.append(int(folder)) 114 | self.img_flag.append(0) 115 | self.samples.append(temp) 116 | else: #val 117 | for folder in os.listdir(self.image_dir): 118 | fdir=self.image_dir+'/'+folder 119 | for files in os.listdir(fdir): 120 | temp=folder+'_'+files 121 | self.img_label.append(int(folder)) 122 | self.img_flag.append(0) 123 | self.samples.append(temp) 124 | 125 | def __len__(self): 126 | return len(self.samples) 127 | 128 | def __getitem__(self, idx): 129 | 130 | temp=self.samples[idx] # folder_files 131 | # print(temp) 132 | if self.img_flag[idx]==1: 133 | foldername='gen_0000' 134 | filename=temp[9:] 135 | else: 136 | foldername=temp[:4] 137 | filename=temp[5:] 138 | img=default_loader(self.image_dir +'/'+foldername+'/'+filename) 139 | if self.train_val=='train_new': 140 | result = {'img': data_transforms['train'](img), 'label': self.img_label[idx], 'flag':self.img_flag[idx]} # flag=0 for ture data and 0 for generated data 141 | else: 142 | result = {'img': data_transforms['val'](img), 'label': self.img_label[idx], 'flag':self.img_flag[idx]} 143 | return result 144 | 145 | class LSROloss(nn.Module): 146 | def __init__(self): # change target to range(0,750) 147 | super(LSROloss,self).__init__() 148 | #input means the prediction score(torch Variable) 32*752,target means the corresponding label, 149 | def forward(self,input,target,flg): # while flg means the flag(=0 for true data and 1 for generated data) batchsize*1 150 | # print(type(input)) 151 | if input.dim()>2: # N defines the number of images, C defines channels, K class in total 152 | input=input.view(input.size(0),input.size(1),-1) # N,C,H,W => N,C,H*W 153 | input=input.transpose(1,2) # N,C,H*W => N,H*W,C 154 | input=input.contiguous().view(-1,input.size(2)) # N,H*W,C => N*H*W,C 155 | 156 | # normalize input 157 | maxRow, _ = torch.max(input.data, 1) # outputs.data return the index of the biggest value in each row 158 | maxRow=maxRow.unsqueeze(1) 159 | input.data=input.data-maxRow 160 | 161 | target=target.view(-1,1) # batchsize*1 162 | flg=flg.view(-1,1) 163 | #len=flg.size()[0] 164 | flos=F.log_softmax(input) # N*K? batchsize*751 165 | flos=torch.sum(flos,1)/flos.size(1) # N*1 get average gan loss 166 | logpt=F.log_softmax(input) # size: batchsize*751 167 | # print(logpt) 168 | logpt=logpt.gather(1,target) # here is a problem 169 | logpt=logpt.view(-1) # N*1 original loss 170 | flg=flg.view(-1) 171 | flg=flg.type(torch.cuda.FloatTensor) 172 | loss=-1*logpt*(1-flg)-flos*flg 173 | return loss.mean() 174 | 175 | dataloaders={} 176 | dataloaders['train'] = DataLoader(dcganDataset('train_new',data_transforms['train']), batch_size=opt.batchsize, 177 | shuffle=True, num_workers=8) 178 | dataloaders['val'] = DataLoader(dcganDataset('val_new',data_transforms['val']), batch_size=opt.batchsize, 179 | shuffle=True, num_workers=8) 180 | 181 | dataset_sizes={} 182 | dataset_train_dir=os.path.join(data_dir,'train_new') 183 | dataset_val_dir=os.path.join(data_dir,'val_new') 184 | dataset_sizes['train']=sum(len(os.listdir(os.path.join(dataset_train_dir,i))) for i in os.listdir(dataset_train_dir)) 185 | dataset_sizes['val']=sum(len(os.listdir(os.path.join(dataset_val_dir,i))) for i in os.listdir(dataset_val_dir)) 186 | 187 | print(dataset_sizes['train']) 188 | print(dataset_sizes['val']) 189 | 190 | #class_names={} 191 | #class_names['train']=len(os.listdir(dataset_train_dir)) 192 | #class_names['val']=len(os.listdir(dataset_val_dir)) 193 | use_gpu = torch.cuda.is_available() 194 | 195 | ###################################################################### 196 | # Training the model 197 | # ------------------ 198 | 199 | y_loss = {} # loss history 200 | y_loss['train'] = [] 201 | y_loss['val'] = [] 202 | y_err = {} 203 | y_err['train'] = [] 204 | y_err['val'] = [] 205 | 206 | def train_model(model, criterion, optimizer, scheduler, num_epochs=25): 207 | since = time.time() 208 | best_model_wts = model.state_dict() 209 | best_acc = 0.0 210 | 211 | for epoch in range(num_epochs): 212 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 213 | print('-' * 10) 214 | 215 | # Each epoch has a training and validation phase 216 | for phase in ['train', 'val']: 217 | if phase == 'train': 218 | scheduler.step() 219 | model.train(True) # Set model to training mode 220 | else: 221 | model.train(False) # Set model to evaluate mode 222 | 223 | running_loss = 0.0 224 | running_corrects = 0 225 | 226 | # Iterate over data. 227 | for data in dataloaders[phase]: 228 | # get the inputs 229 | inputs=data['img'] 230 | labels=data['label'] 231 | flags= data['flag'] 232 | 233 | if use_gpu: 234 | inputs = Variable(inputs.cuda()) 235 | labels = Variable(labels.cuda()) 236 | flags=Variable(flags.cuda()) 237 | else: 238 | inputs, labels,flags = Variable(inputs), Variable(labels), Variable(flags) 239 | 240 | # zero the parameter gradients 241 | optimizer.zero_grad() 242 | 243 | # forward 244 | outputs = model(inputs) 245 | _, preds = torch.max(outputs.data, 1) # outputs.data return the index of the biggest value in each row 246 | loss = criterion(outputs,labels,flags) 247 | 248 | # backward + optimize only if in training phase 249 | if phase == 'train': 250 | loss.backward() 251 | optimizer.step() 252 | 253 | # statistics 254 | running_loss += loss.data[0] 255 | 256 | for temp in range(flags.size()[0]): 257 | if flags.data[temp]==1: 258 | preds[temp]=-1 259 | 260 | running_corrects += torch.sum(preds == labels.data) 261 | # print('running_corrects: '+str(running_corrects)) 262 | 263 | epoch_loss = running_loss / dataset_sizes[phase] 264 | #epoch_acc = running_corrects / dataset_sizes[phase] 265 | if phase =='train': 266 | # epoch_acc = running_corrects / (dataset_sizes[phase]-4992) # 4992 generated image in total 267 | epoch_acc = running_corrects / (dataset_sizes[phase]-generated_image_size) # 4992 generated image in total 268 | else: 269 | epoch_acc = running_corrects / dataset_sizes[phase] 270 | 271 | print('{} Loss: {:.4f} Acc: {:.4f}'.format( 272 | phase, epoch_loss, epoch_acc)) 273 | y_loss[phase].append(epoch_loss) 274 | y_err[phase].append(1.0-epoch_acc) 275 | # deep copy the model 276 | if phase == 'val': 277 | if epoch_acc>best_acc: 278 | best_acc=epoch_acc 279 | best_model_wts = model.state_dict() 280 | if epoch>=40: 281 | save_network(model, epoch) 282 | # draw_curve(epoch) 283 | 284 | print() 285 | 286 | time_elapsed = time.time() - since 287 | print('Training complete in {:.0f}m {:.0f}s'.format( 288 | time_elapsed // 60, time_elapsed % 60)) 289 | print('Best val Acc: {:4f}'.format(best_acc)) 290 | 291 | # load best model weights 292 | model.load_state_dict(best_model_wts) 293 | save_network(model, 'best') 294 | return model 295 | 296 | 297 | ###################################################################### 298 | # Save model 299 | #--------------------------- 300 | def save_network(network, epoch_label): 301 | save_filename = 'net_%s.pth'% epoch_label 302 | save_path = os.path.join('./model',name,save_filename) 303 | torch.save(network.state_dict(), save_path) 304 | # this step is important, or error occurs "runtimeError: tensors are on different GPUs" 305 | # if torch.cuda.is_available: 306 | # network.cuda(gpu_ids[0]) 307 | #if torch.cuda.is_available: 308 | # network=nn.DataParallel(network,device_ids=[0,1,2]) # multi-GPU 309 | 310 | 311 | #print('------------'+str(len(clas_names))+'--------------') 312 | if opt.use_dense: 313 | #print(len(class_names['train'])) 314 | model = ft_net_dense(751) # 751 class for training data in market 1501 in total 315 | else: 316 | model = ft_net(751) 317 | 318 | if use_gpu: 319 | model = model.cuda() 320 | criterion = LSROloss() 321 | 322 | ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.classifier.parameters() )) 323 | base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) 324 | 325 | # Observe that all parameters are being optimized 326 | optimizer_ft = optim.SGD([ 327 | {'params': base_params, 'lr': 0.01}, 328 | {'params': model.model.fc.parameters(), 'lr': 0.05}, 329 | {'params': model.classifier.parameters(), 'lr': 0.05} 330 | ], momentum=0.9, weight_decay=5e-4, nesterov=True) 331 | 332 | model=nn.DataParallel(model,device_ids=[0,1,2]) # multi-GPU 333 | 334 | # Decay LR by a factor of 0.1 every 40 epochs 335 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=40, gamma=0.1) 336 | 337 | dir_name = os.path.join('./model',name) 338 | if not os.path.isdir(dir_name): 339 | os.mkdir(dir_name) 340 | 341 | # save opts 342 | with open('%s/opts.json'%dir_name,'w') as fp: 343 | json.dump(vars(opt), fp, indent=1) 344 | 345 | model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 346 | num_epochs=130) 347 | -------------------------------------------------------------------------------- /DCGAN-tensorflow/model.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | import time 4 | import math 5 | from glob import glob 6 | import tensorflow as tf 7 | import numpy as np 8 | from six.moves import xrange 9 | 10 | from ops import * 11 | from utils import * 12 | 13 | def conv_out_size_same(size, stride): 14 | return int(math.ceil(float(size) / float(stride))) 15 | 16 | class DCGAN(object): 17 | def __init__(self, sess, input_height=108, input_width=108, crop=True, 18 | batch_size=64, sample_num = 64, output_height=64, output_width=64, 19 | y_dim=None, z_dim=100, gf_dim=64, df_dim=64, 20 | gfc_dim=1024, dfc_dim=1024, c_dim=3, dataset_name='default', 21 | input_fname_pattern='*.jpg', checkpoint_dir=None, sample_dir=None): 22 | """ 23 | 24 | Args: 25 | sess: TensorFlow session 26 | batch_size: The size of batch. Should be specified before training. 27 | y_dim: (optional) Dimension of dim for y. [None] 28 | z_dim: (optional) Dimension of dim for Z. [100] 29 | gf_dim: (optional) Dimension of gen filters in first conv layer. [64] 30 | df_dim: (optional) Dimension of discrim filters in first conv layer. [64] 31 | gfc_dim: (optional) Dimension of gen units for for fully connected layer. [1024] 32 | dfc_dim: (optional) Dimension of discrim units for fully connected layer. [1024] 33 | c_dim: (optional) Dimension of image color. For grayscale input, set to 1. [3] 34 | """ 35 | self.sess = sess 36 | self.crop = crop 37 | 38 | self.batch_size = batch_size 39 | self.sample_num = sample_num 40 | 41 | self.input_height = input_height 42 | self.input_width = input_width 43 | self.output_height = output_height 44 | self.output_width = output_width 45 | 46 | self.y_dim = y_dim 47 | self.z_dim = z_dim 48 | 49 | self.gf_dim = gf_dim 50 | self.df_dim = df_dim 51 | 52 | self.gfc_dim = gfc_dim 53 | self.dfc_dim = dfc_dim 54 | 55 | # batch normalization : deals with poor initialization helps gradient flow 56 | self.d_bn1 = batch_norm(name='d_bn1') 57 | self.d_bn2 = batch_norm(name='d_bn2') 58 | 59 | if not self.y_dim: 60 | self.d_bn3 = batch_norm(name='d_bn3') 61 | 62 | self.g_bn0 = batch_norm(name='g_bn0') 63 | self.g_bn1 = batch_norm(name='g_bn1') 64 | self.g_bn2 = batch_norm(name='g_bn2') 65 | 66 | if not self.y_dim: 67 | self.g_bn3 = batch_norm(name='g_bn3') 68 | 69 | self.dataset_name = dataset_name 70 | self.input_fname_pattern = input_fname_pattern 71 | self.checkpoint_dir = checkpoint_dir 72 | 73 | if self.dataset_name == 'mnist': 74 | self.data_X, self.data_y = self.load_mnist() 75 | self.c_dim = self.data_X[0].shape[-1] 76 | else: 77 | self.data = glob(os.path.join("./data", self.dataset_name, self.input_fname_pattern)) 78 | imreadImg = imread(self.data[0]); 79 | if len(imreadImg.shape) >= 3: #check if image is a non-grayscale image by checking channel number 80 | self.c_dim = imread(self.data[0]).shape[-1] 81 | else: 82 | self.c_dim = 1 83 | 84 | self.grayscale = (self.c_dim == 1) 85 | 86 | self.build_model() 87 | 88 | def build_model(self): 89 | if self.y_dim: 90 | self.y = tf.placeholder(tf.float32, [self.batch_size, self.y_dim], name='y') 91 | else: 92 | self.y = None 93 | 94 | if self.crop: 95 | image_dims = [self.output_height, self.output_width, self.c_dim] 96 | else: 97 | image_dims = [self.input_height, self.input_width, self.c_dim] 98 | 99 | self.inputs = tf.placeholder( 100 | tf.float32, [self.batch_size] + image_dims, name='real_images') 101 | 102 | inputs = self.inputs 103 | 104 | self.z = tf.placeholder( 105 | tf.float32, [None, self.z_dim], name='z') 106 | self.z_sum = histogram_summary("z", self.z) 107 | 108 | self.G = self.generator(self.z, self.y) 109 | self.D, self.D_logits = self.discriminator(inputs, self.y, reuse=False) 110 | self.sampler = self.sampler(self.z, self.y) 111 | self.D_, self.D_logits_ = self.discriminator(self.G, self.y, reuse=True) 112 | 113 | self.d_sum = histogram_summary("d", self.D) 114 | self.d__sum = histogram_summary("d_", self.D_) 115 | self.G_sum = image_summary("G", self.G) 116 | 117 | def sigmoid_cross_entropy_with_logits(x, y): 118 | try: 119 | return tf.nn.sigmoid_cross_entropy_with_logits(logits=x, labels=y) 120 | except: 121 | return tf.nn.sigmoid_cross_entropy_with_logits(logits=x, targets=y) 122 | 123 | self.d_loss_real = tf.reduce_mean( 124 | sigmoid_cross_entropy_with_logits(self.D_logits, tf.ones_like(self.D))) 125 | self.d_loss_fake = tf.reduce_mean( 126 | sigmoid_cross_entropy_with_logits(self.D_logits_, tf.zeros_like(self.D_))) 127 | self.g_loss = tf.reduce_mean( 128 | sigmoid_cross_entropy_with_logits(self.D_logits_, tf.ones_like(self.D_))) 129 | 130 | self.d_loss_real_sum = scalar_summary("d_loss_real", self.d_loss_real) 131 | self.d_loss_fake_sum = scalar_summary("d_loss_fake", self.d_loss_fake) 132 | 133 | self.d_loss = self.d_loss_real + self.d_loss_fake 134 | 135 | self.g_loss_sum = scalar_summary("g_loss", self.g_loss) 136 | self.d_loss_sum = scalar_summary("d_loss", self.d_loss) 137 | 138 | t_vars = tf.trainable_variables() 139 | 140 | self.d_vars = [var for var in t_vars if 'd_' in var.name] 141 | self.g_vars = [var for var in t_vars if 'g_' in var.name] 142 | 143 | self.saver = tf.train.Saver() 144 | 145 | def train(self, config): 146 | d_optim = tf.train.AdamOptimizer(config.learning_rate, beta1=config.beta1) \ 147 | .minimize(self.d_loss, var_list=self.d_vars) 148 | g_optim = tf.train.AdamOptimizer(config.learning_rate, beta1=config.beta1) \ 149 | .minimize(self.g_loss, var_list=self.g_vars) 150 | try: 151 | tf.global_variables_initializer().run() 152 | except: 153 | tf.initialize_all_variables().run() 154 | 155 | self.g_sum = merge_summary([self.z_sum, self.d__sum, 156 | self.G_sum, self.d_loss_fake_sum, self.g_loss_sum]) 157 | self.d_sum = merge_summary( 158 | [self.z_sum, self.d_sum, self.d_loss_real_sum, self.d_loss_sum]) 159 | self.writer = SummaryWriter("./logs", self.sess.graph) 160 | 161 | sample_z = np.random.uniform(-1, 1, size=(self.sample_num , self.z_dim)) 162 | 163 | if config.dataset == 'mnist': 164 | sample_inputs = self.data_X[0:self.sample_num] 165 | sample_labels = self.data_y[0:self.sample_num] 166 | else: 167 | sample_files = self.data[0:self.sample_num] 168 | sample = [ 169 | get_image(sample_file, 170 | input_height=self.input_height, 171 | input_width=self.input_width, 172 | resize_height=self.output_height, 173 | resize_width=self.output_width, 174 | crop=self.crop, 175 | grayscale=self.grayscale) for sample_file in sample_files] 176 | if (self.grayscale): 177 | sample_inputs = np.array(sample).astype(np.float32)[:, :, :, None] 178 | else: 179 | sample_inputs = np.array(sample).astype(np.float32) 180 | 181 | counter = 1 182 | start_time = time.time() 183 | could_load, checkpoint_counter = self.load(self.checkpoint_dir) 184 | if could_load: 185 | counter = checkpoint_counter 186 | print(" [*] Load SUCCESS") 187 | else: 188 | print(" [!] Load failed...") 189 | 190 | for epoch in xrange(config.epoch): 191 | if config.dataset == 'mnist': 192 | batch_idxs = min(len(self.data_X), config.train_size) // config.batch_size 193 | else: 194 | self.data = glob(os.path.join( 195 | "./data", config.dataset, self.input_fname_pattern)) 196 | batch_idxs = min(len(self.data), config.train_size) // config.batch_size 197 | 198 | for idx in xrange(0, batch_idxs): 199 | if config.dataset == 'mnist': 200 | batch_images = self.data_X[idx*config.batch_size:(idx+1)*config.batch_size] 201 | batch_labels = self.data_y[idx*config.batch_size:(idx+1)*config.batch_size] 202 | else: 203 | batch_files = self.data[idx*config.batch_size:(idx+1)*config.batch_size] 204 | batch = [ 205 | get_image(batch_file, 206 | input_height=self.input_height, 207 | input_width=self.input_width, 208 | resize_height=self.output_height, 209 | resize_width=self.output_width, 210 | crop=self.crop, 211 | grayscale=self.grayscale) for batch_file in batch_files] 212 | if self.grayscale: 213 | batch_images = np.array(batch).astype(np.float32)[:, :, :, None] 214 | else: 215 | batch_images = np.array(batch).astype(np.float32) 216 | 217 | batch_z = np.random.uniform(-1, 1, [config.batch_size, self.z_dim]) \ 218 | .astype(np.float32) 219 | 220 | if config.dataset == 'mnist': 221 | # Update D network 222 | _, summary_str = self.sess.run([d_optim, self.d_sum], 223 | feed_dict={ 224 | self.inputs: batch_images, 225 | self.z: batch_z, 226 | self.y:batch_labels, 227 | }) 228 | self.writer.add_summary(summary_str, counter) 229 | 230 | # Update G network 231 | _, summary_str = self.sess.run([g_optim, self.g_sum], 232 | feed_dict={ 233 | self.z: batch_z, 234 | self.y:batch_labels, 235 | }) 236 | self.writer.add_summary(summary_str, counter) 237 | 238 | # Run g_optim twice to make sure that d_loss does not go to zero (different from paper) 239 | _, summary_str = self.sess.run([g_optim, self.g_sum], 240 | feed_dict={ self.z: batch_z, self.y:batch_labels }) 241 | self.writer.add_summary(summary_str, counter) 242 | 243 | errD_fake = self.d_loss_fake.eval({ 244 | self.z: batch_z, 245 | self.y:batch_labels 246 | }) 247 | errD_real = self.d_loss_real.eval({ 248 | self.inputs: batch_images, 249 | self.y:batch_labels 250 | }) 251 | errG = self.g_loss.eval({ 252 | self.z: batch_z, 253 | self.y: batch_labels 254 | }) 255 | else: 256 | # Update D network 257 | _, summary_str = self.sess.run([d_optim, self.d_sum], 258 | feed_dict={ self.inputs: batch_images, self.z: batch_z }) 259 | self.writer.add_summary(summary_str, counter) 260 | 261 | # Update G network 262 | _, summary_str = self.sess.run([g_optim, self.g_sum], 263 | feed_dict={ self.z: batch_z }) 264 | self.writer.add_summary(summary_str, counter) 265 | 266 | # Run g_optim twice to make sure that d_loss does not go to zero (different from paper) 267 | _, summary_str = self.sess.run([g_optim, self.g_sum], 268 | feed_dict={ self.z: batch_z }) 269 | self.writer.add_summary(summary_str, counter) 270 | 271 | errD_fake = self.d_loss_fake.eval({ self.z: batch_z }) 272 | errD_real = self.d_loss_real.eval({ self.inputs: batch_images }) 273 | errG = self.g_loss.eval({self.z: batch_z}) 274 | 275 | counter += 1 276 | print("Epoch: [%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \ 277 | % (epoch, idx, batch_idxs, 278 | time.time() - start_time, errD_fake+errD_real, errG)) 279 | 280 | if np.mod(counter, 100) == 1: 281 | if config.dataset == 'mnist': 282 | samples, d_loss, g_loss = self.sess.run( 283 | [self.sampler, self.d_loss, self.g_loss], 284 | feed_dict={ 285 | self.z: sample_z, 286 | self.inputs: sample_inputs, 287 | self.y:sample_labels, 288 | } 289 | ) 290 | save_images(samples, image_manifold_size(samples.shape[0]), 291 | './{}/train_{:02d}_{:04d}.png'.format(config.sample_dir, epoch, idx)) 292 | print("[Sample] d_loss: %.8f, g_loss: %.8f" % (d_loss, g_loss)) 293 | else: 294 | try: 295 | samples, d_loss, g_loss = self.sess.run( 296 | [self.sampler, self.d_loss, self.g_loss], 297 | feed_dict={ 298 | self.z: sample_z, 299 | self.inputs: sample_inputs, 300 | }, 301 | ) 302 | save_images(samples, image_manifold_size(samples.shape[0]), 303 | './{}/train_{:02d}_{:04d}.png'.format(config.sample_dir, epoch, idx)) 304 | print("[Sample] d_loss: %.8f, g_loss: %.8f" % (d_loss, g_loss)) 305 | except: 306 | print("one pic error!...") 307 | 308 | if np.mod(counter, 500) == 2: 309 | self.save(config.checkpoint_dir, counter) 310 | 311 | def discriminator(self, image, y=None, reuse=False): 312 | with tf.variable_scope("discriminator") as scope: 313 | if reuse: 314 | scope.reuse_variables() 315 | 316 | if not self.y_dim: 317 | h0 = lrelu(conv2d(image, self.df_dim, name='d_h0_conv')) 318 | h1 = lrelu(self.d_bn1(conv2d(h0, self.df_dim*2, name='d_h1_conv'))) 319 | h2 = lrelu(self.d_bn2(conv2d(h1, self.df_dim*4, name='d_h2_conv'))) 320 | h3 = lrelu(self.d_bn3(conv2d(h2, self.df_dim*8, name='d_h3_conv'))) 321 | h4 = linear(tf.reshape(h3, [self.batch_size, -1]), 1, 'd_h4_lin') 322 | 323 | return tf.nn.sigmoid(h4), h4 324 | else: 325 | yb = tf.reshape(y, [self.batch_size, 1, 1, self.y_dim]) 326 | x = conv_cond_concat(image, yb) 327 | 328 | h0 = lrelu(conv2d(x, self.c_dim + self.y_dim, name='d_h0_conv')) 329 | h0 = conv_cond_concat(h0, yb) 330 | 331 | h1 = lrelu(self.d_bn1(conv2d(h0, self.df_dim + self.y_dim, name='d_h1_conv'))) 332 | h1 = tf.reshape(h1, [self.batch_size, -1]) 333 | h1 = concat([h1, y], 1) 334 | 335 | h2 = lrelu(self.d_bn2(linear(h1, self.dfc_dim, 'd_h2_lin'))) 336 | h2 = concat([h2, y], 1) 337 | 338 | h3 = linear(h2, 1, 'd_h3_lin') 339 | 340 | return tf.nn.sigmoid(h3), h3 341 | 342 | def generator(self, z, y=None): 343 | with tf.variable_scope("generator") as scope: 344 | if not self.y_dim: 345 | s_h, s_w = self.output_height, self.output_width 346 | s_h2, s_w2 = conv_out_size_same(s_h, 2), conv_out_size_same(s_w, 2) 347 | s_h4, s_w4 = conv_out_size_same(s_h2, 2), conv_out_size_same(s_w2, 2) 348 | s_h8, s_w8 = conv_out_size_same(s_h4, 2), conv_out_size_same(s_w4, 2) 349 | s_h16, s_w16 = conv_out_size_same(s_h8, 2), conv_out_size_same(s_w8, 2) 350 | 351 | # project `z` and reshape 352 | self.z_, self.h0_w, self.h0_b = linear( 353 | z, self.gf_dim*8*s_h16*s_w16, 'g_h0_lin', with_w=True) 354 | 355 | self.h0 = tf.reshape( 356 | self.z_, [-1, s_h16, s_w16, self.gf_dim * 8]) 357 | h0 = tf.nn.relu(self.g_bn0(self.h0)) 358 | 359 | self.h1, self.h1_w, self.h1_b = deconv2d( 360 | h0, [self.batch_size, s_h8, s_w8, self.gf_dim*4], name='g_h1', with_w=True) 361 | h1 = tf.nn.relu(self.g_bn1(self.h1)) 362 | 363 | h2, self.h2_w, self.h2_b = deconv2d( 364 | h1, [self.batch_size, s_h4, s_w4, self.gf_dim*2], name='g_h2', with_w=True) 365 | h2 = tf.nn.relu(self.g_bn2(h2)) 366 | 367 | h3, self.h3_w, self.h3_b = deconv2d( 368 | h2, [self.batch_size, s_h2, s_w2, self.gf_dim*1], name='g_h3', with_w=True) 369 | h3 = tf.nn.relu(self.g_bn3(h3)) 370 | 371 | h4, self.h4_w, self.h4_b = deconv2d( 372 | h3, [self.batch_size, s_h, s_w, self.c_dim], name='g_h4', with_w=True) 373 | 374 | return tf.nn.tanh(h4) 375 | else: 376 | s_h, s_w = self.output_height, self.output_width 377 | s_h2, s_h4 = int(s_h/2), int(s_h/4) 378 | s_w2, s_w4 = int(s_w/2), int(s_w/4) 379 | 380 | # yb = tf.expand_dims(tf.expand_dims(y, 1),2) 381 | yb = tf.reshape(y, [self.batch_size, 1, 1, self.y_dim]) 382 | z = concat([z, y], 1) 383 | 384 | h0 = tf.nn.relu( 385 | self.g_bn0(linear(z, self.gfc_dim, 'g_h0_lin'))) 386 | h0 = concat([h0, y], 1) 387 | 388 | h1 = tf.nn.relu(self.g_bn1( 389 | linear(h0, self.gf_dim*2*s_h4*s_w4, 'g_h1_lin'))) 390 | h1 = tf.reshape(h1, [self.batch_size, s_h4, s_w4, self.gf_dim * 2]) 391 | 392 | h1 = conv_cond_concat(h1, yb) 393 | 394 | h2 = tf.nn.relu(self.g_bn2(deconv2d(h1, 395 | [self.batch_size, s_h2, s_w2, self.gf_dim * 2], name='g_h2'))) 396 | h2 = conv_cond_concat(h2, yb) 397 | 398 | return tf.nn.sigmoid( 399 | deconv2d(h2, [self.batch_size, s_h, s_w, self.c_dim], name='g_h3')) 400 | 401 | def sampler(self, z, y=None): 402 | with tf.variable_scope("generator") as scope: 403 | scope.reuse_variables() 404 | 405 | if not self.y_dim: 406 | s_h, s_w = self.output_height, self.output_width 407 | s_h2, s_w2 = conv_out_size_same(s_h, 2), conv_out_size_same(s_w, 2) 408 | s_h4, s_w4 = conv_out_size_same(s_h2, 2), conv_out_size_same(s_w2, 2) 409 | s_h8, s_w8 = conv_out_size_same(s_h4, 2), conv_out_size_same(s_w4, 2) 410 | s_h16, s_w16 = conv_out_size_same(s_h8, 2), conv_out_size_same(s_w8, 2) 411 | 412 | # project `z` and reshape 413 | h0 = tf.reshape( 414 | linear(z, self.gf_dim*8*s_h16*s_w16, 'g_h0_lin'), 415 | [-1, s_h16, s_w16, self.gf_dim * 8]) 416 | h0 = tf.nn.relu(self.g_bn0(h0, train=False)) 417 | 418 | h1 = deconv2d(h0, [self.batch_size, s_h8, s_w8, self.gf_dim*4], name='g_h1') 419 | h1 = tf.nn.relu(self.g_bn1(h1, train=False)) 420 | 421 | h2 = deconv2d(h1, [self.batch_size, s_h4, s_w4, self.gf_dim*2], name='g_h2') 422 | h2 = tf.nn.relu(self.g_bn2(h2, train=False)) 423 | 424 | h3 = deconv2d(h2, [self.batch_size, s_h2, s_w2, self.gf_dim*1], name='g_h3') 425 | h3 = tf.nn.relu(self.g_bn3(h3, train=False)) 426 | 427 | h4 = deconv2d(h3, [self.batch_size, s_h, s_w, self.c_dim], name='g_h4') 428 | 429 | return tf.nn.tanh(h4) 430 | else: 431 | s_h, s_w = self.output_height, self.output_width 432 | s_h2, s_h4 = int(s_h/2), int(s_h/4) 433 | s_w2, s_w4 = int(s_w/2), int(s_w/4) 434 | 435 | # yb = tf.reshape(y, [-1, 1, 1, self.y_dim]) 436 | yb = tf.reshape(y, [self.batch_size, 1, 1, self.y_dim]) 437 | z = concat([z, y], 1) 438 | 439 | h0 = tf.nn.relu(self.g_bn0(linear(z, self.gfc_dim, 'g_h0_lin'), train=False)) 440 | h0 = concat([h0, y], 1) 441 | 442 | h1 = tf.nn.relu(self.g_bn1( 443 | linear(h0, self.gf_dim*2*s_h4*s_w4, 'g_h1_lin'), train=False)) 444 | h1 = tf.reshape(h1, [self.batch_size, s_h4, s_w4, self.gf_dim * 2]) 445 | h1 = conv_cond_concat(h1, yb) 446 | 447 | h2 = tf.nn.relu(self.g_bn2( 448 | deconv2d(h1, [self.batch_size, s_h2, s_w2, self.gf_dim * 2], name='g_h2'), train=False)) 449 | h2 = conv_cond_concat(h2, yb) 450 | 451 | return tf.nn.sigmoid(deconv2d(h2, [self.batch_size, s_h, s_w, self.c_dim], name='g_h3')) 452 | 453 | def load_mnist(self): 454 | data_dir = os.path.join("./data", self.dataset_name) 455 | 456 | fd = open(os.path.join(data_dir,'train-images-idx3-ubyte')) 457 | loaded = np.fromfile(file=fd,dtype=np.uint8) 458 | trX = loaded[16:].reshape((60000,28,28,1)).astype(np.float) 459 | 460 | fd = open(os.path.join(data_dir,'train-labels-idx1-ubyte')) 461 | loaded = np.fromfile(file=fd,dtype=np.uint8) 462 | trY = loaded[8:].reshape((60000)).astype(np.float) 463 | 464 | fd = open(os.path.join(data_dir,'t10k-images-idx3-ubyte')) 465 | loaded = np.fromfile(file=fd,dtype=np.uint8) 466 | teX = loaded[16:].reshape((10000,28,28,1)).astype(np.float) 467 | 468 | fd = open(os.path.join(data_dir,'t10k-labels-idx1-ubyte')) 469 | loaded = np.fromfile(file=fd,dtype=np.uint8) 470 | teY = loaded[8:].reshape((10000)).astype(np.float) 471 | 472 | trY = np.asarray(trY) 473 | teY = np.asarray(teY) 474 | 475 | X = np.concatenate((trX, teX), axis=0) 476 | y = np.concatenate((trY, teY), axis=0).astype(np.int) 477 | 478 | seed = 547 479 | np.random.seed(seed) 480 | np.random.shuffle(X) 481 | np.random.seed(seed) 482 | np.random.shuffle(y) 483 | 484 | y_vec = np.zeros((len(y), self.y_dim), dtype=np.float) 485 | for i, label in enumerate(y): 486 | y_vec[i,y[i]] = 1.0 487 | 488 | return X/255.,y_vec 489 | 490 | @property 491 | def model_dir(self): 492 | return "{}_{}_{}_{}".format( 493 | self.dataset_name, self.batch_size, 494 | self.output_height, self.output_width) 495 | 496 | def save(self, checkpoint_dir, step): 497 | model_name = "DCGAN.model" 498 | checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir) 499 | 500 | if not os.path.exists(checkpoint_dir): 501 | os.makedirs(checkpoint_dir) 502 | 503 | self.saver.save(self.sess, 504 | os.path.join(checkpoint_dir, model_name), 505 | global_step=step) 506 | 507 | def load(self, checkpoint_dir): 508 | import re 509 | print(" [*] Reading checkpoints...") 510 | checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir) 511 | 512 | ckpt = tf.train.get_checkpoint_state(checkpoint_dir) 513 | if ckpt and ckpt.model_checkpoint_path: 514 | ckpt_name = os.path.basename(ckpt.model_checkpoint_path) 515 | self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name)) 516 | counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0)) 517 | print(" [*] Success to read {}".format(ckpt_name)) 518 | return True, counter 519 | else: 520 | print(" [*] Failed to find a checkpoint") 521 | return False, 0 522 | --------------------------------------------------------------------------------