diff --git a/.gitignore b/.gitignore index 3b35c41..4a89ee1 100644 --- a/.gitignore +++ b/.gitignore @@ -66,7 +66,7 @@ clearml.conf # etc -datafolder/ -doc/ -net/ -test_sample/ +# datafolder/ +# doc/ +# net/ +# test_sample/ diff --git a/datafolder/folder.py b/datafolder/folder.py new file mode 100644 index 0000000..2948ca6 --- /dev/null +++ b/datafolder/folder.py @@ -0,0 +1,140 @@ +import os +from PIL import Image +import torch +from torch.utils import data +import numpy as np +from torchvision import transforms as T +from .reid_dataset import import_MarketDuke_nodistractors +from .reid_dataset import import_Market1501Attribute_binary +from .reid_dataset import import_DukeMTMCAttribute_binary + + +class Train_Dataset(data.Dataset): + + def __init__(self, data_dir, dataset_name, transforms=None, train_val='train' ): + + train, query, gallery = import_MarketDuke_nodistractors(data_dir, dataset_name) + + if dataset_name == 'Market-1501': + train_attr, test_attr, self.label = import_Market1501Attribute_binary(data_dir) + elif dataset_name == 'DukeMTMC-reID': + train_attr, test_attr, self.label = import_DukeMTMCAttribute_binary(data_dir) + else: + print('Input should only be Market1501 or DukeMTMC') + + self.num_ids = len(train['ids']) + self.num_labels = len(self.label) + + # distribution:每个属性的正样本占比 + distribution = np.zeros(self.num_labels) + for k, v in train_attr.items(): + distribution += np.array(v) + self.distribution = distribution / len(train_attr) + + if train_val == 'train': + self.train_data = train['data'] + self.train_ids = train['ids'] + self.train_attr = train_attr + elif train_val == 'query': + self.train_data = query['data'] + self.train_ids = query['ids'] + self.train_attr = test_attr + elif train_val == 'gallery': + self.train_data = gallery['data'] + self.train_ids = gallery['ids'] + self.train_attr = test_attr + else: + print('Input should only be train or val') + + self.num_ids = len(self.train_ids) + + if transforms is None: + if train_val == 'train': + self.transforms = T.Compose([ + T.Resize(size=(288, 144)), + T.RandomHorizontalFlip(), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + else: + self.transforms = T.Compose([ + T.Resize(size=(288, 144)), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + def __getitem__(self, index): + ''' + 一次返回一张图片的数据 + ''' + img_path = self.train_data[index][0] + i = self.train_data[index][1] + id = self.train_data[index][2] + cam = self.train_data[index][3] + label = np.asarray(self.train_attr[id]) + data = Image.open(img_path) + data = self.transforms(data) + name = self.train_data[index][4] + return data, i, label, id, cam, name + + def __len__(self): + return len(self.train_data) + + def num_label(self): + return self.num_labels + + def num_id(self): + return self.num_ids + + def labels(self): + return self.label + + + +class Test_Dataset(data.Dataset): + def __init__(self, data_dir, dataset_name, transforms=None, query_gallery='query' ): + train, query, gallery = import_MarketDuke_nodistractors(data_dir, dataset_name) + + if dataset_name == 'Market-1501': + self.train_attr, self.test_attr, self.label = import_Market1501Attribute_binary(data_dir) + elif dataset_name == 'DukeMTMC-reID': + self.train_attr, self.test_attr, self.label = import_DukeMTMCAttribute_binary(data_dir) + else: + print('Input should only be Market1501 or DukeMTMC') + + if query_gallery == 'query': + self.test_data = query['data'] + self.test_ids = query['ids'] + elif query_gallery == 'gallery': + self.test_data = gallery['data'] + self.test_ids = gallery['ids'] + elif query_gallery == 'all': + self.test_data = gallery['data'] + query['data'] + self.test_ids = gallery['ids'] + else: + print('Input shoud only be query or gallery;') + + if transforms is None: + self.transforms = T.Compose([ + T.Resize(size=(288, 144)), + T.ToTensor(), + T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + def __getitem__(self, index): + ''' + 一次返回一张图片的数据 + ''' + img_path = self.test_data[index][0] + id = self.test_data[index][2] + label = np.asarray(self.test_attr[id]) + data = Image.open(img_path) + data = self.transforms(data) + name = self.test_data[index][4] + return data, label, id, name + + def __len__(self): + return len(self.test_data) + + def labels(self): + return self.label \ No newline at end of file diff --git a/datafolder/reid_dataset/__init__.py b/datafolder/reid_dataset/__init__.py new file mode 100644 index 0000000..f8f7392 --- /dev/null +++ b/datafolder/reid_dataset/__init__.py @@ -0,0 +1,17 @@ +from .reiddataset_downloader import reiddataset_downloader +from .reiddataset_downloader import reiddataset_downloader_all +from .import_VIPeR import import_VIPeR +from .import_CUHK01 import import_CUHK01 +from .import_CUHK03 import import_CUHK03 +from .import_Market1501 import import_Market1501 +from .import_Market1501Attribute import import_Market1501Attribute +from .import_Market1501Attribute import import_Market1501Attribute_binary +from .import_DukeMTMC import import_DukeMTMC +from .import_DukeMTMCAttribute import import_DukeMTMCAttribute +from .import_DukeMTMCAttribute import import_DukeMTMCAttribute_binary +from .import_MarketDuke import import_MarketDuke +from .import_MarketDuke_nodistractors import import_MarketDuke_nodistractors +from .pytorch_prepare import pytorch_prepare +from .pytorch_prepare import pytorch_prepare_all +from .marketduke_to_hdf5 import marketduke_to_hdf5 +from .cuhk03_to_image import cuhk03_to_image diff --git a/datafolder/reid_dataset/cuhk03_to_image.py b/datafolder/reid_dataset/cuhk03_to_image.py new file mode 100644 index 0000000..12c7d90 --- /dev/null +++ b/datafolder/reid_dataset/cuhk03_to_image.py @@ -0,0 +1,47 @@ +import warnings +warnings.filterwarnings('ignore','.*conversion.*') + +import os +import zipfile +import shutil +import requests +import h5py +import numpy as np +from PIL import Image +import argparse + +def cuhk03_to_image(CUHK03_dir): + + f = h5py.File(os.path.join(CUHK03_dir,'cuhk-03.mat')) + + detected_labeled = ['detected','labeled'] + print('converting') + for data_type in detected_labeled: + + datatype_dir = os.path.join(CUHK03_dir, data_type) + if not os.path.exists(datatype_dir): + os.makedirs(datatype_dir) + + for campair in range(len(f[data_type][0])): + campair_dir = os.path.join(datatype_dir,'P%d'%(campair+1)) + cam1_dir = os.path.join(campair_dir,'cam1') + cam2_dir = os.path.join(campair_dir,'cam2') + + if not os.path.exists(campair_dir): + os.makedirs(campair_dir) + if not os.path.exists(cam1_dir): + os.makedirs(cam1_dir) + if not os.path.exists(cam2_dir): + os.makedirs(cam2_dir) + + for img_no in range(f[f[data_type][0][campair]].shape[0]): + if img_no < 5: + cam_dir = 'cam1' + else: + cam_dir = 'cam2' + for person_id in range(f[f[data_type][0][campair]].shape[1]): + img = np.array(f[f[f[data_type][0][campair]][img_no][person_id]]) + if img.shape[0] !=2: + img = np.transpose(img, (2,1,0)) + im = Image.fromarray(img) + im.save(os.path.join(campair_dir, cam_dir, "%d-%d.jpg"%(person_id+1,img_no+1))) \ No newline at end of file diff --git a/datafolder/reid_dataset/gdrive_downloader.py b/datafolder/reid_dataset/gdrive_downloader.py new file mode 100644 index 0000000..8fa6276 --- /dev/null +++ b/datafolder/reid_dataset/gdrive_downloader.py @@ -0,0 +1,37 @@ +import requests + +def gdrive_downloader(destination, id): + URL = "https://docs.google.com/uc?export=download" + + session = requests.Session() + + response = session.get(URL, params = { 'id' : id }, stream = True) + token = get_confirm_token(response) + + if token: + params = { 'id' : id, 'confirm' : token } + response = session.get(URL, params = params, stream = True) + + save_response_content(response, destination) + +def get_confirm_token(response): + for key, value in response.cookies.items(): + if key.startswith('download_warning'): + return value + + return None + +def save_response_content(response, destination): + CHUNK_SIZE = 32768 + + with open(destination, "wb") as f: + for chunk in response.iter_content(CHUNK_SIZE): + if chunk: # filter out keep-alive new chunks + f.write(chunk) + +if __name__ == "__main__": + var = raw_input("Please enter public file id : ") + file_id = str(var) + name = raw_input("Please enter name with extension : ") + destination = str(name) + gdrive_downloader(file_id, destination) diff --git a/datafolder/reid_dataset/import_CUHK01.py b/datafolder/reid_dataset/import_CUHK01.py new file mode 100644 index 0000000..b87c86c --- /dev/null +++ b/datafolder/reid_dataset/import_CUHK01.py @@ -0,0 +1,22 @@ +import os +from .reiddataset_downloader import * +def import_CUHK01(dataset_dir): + cuhk01_dir = os.path.join(dataset_dir,'CUHK01') + + if not os.path.exists(cuhk01_dir): + print('Please Download the CUHK01 Dataset') + + file_list=os.listdir(cuhk01_dir) + name_dict={} + for name in file_list: + if name[-3:]=='png': + id = name[:4] + if id not in name_dict: + name_dict[id]=[] + name_dict[id].append([]) + name_dict[id].append([]) + if int(name[-7:-4])<3: + name_dict[id][0].append(os.path.join(cuhk01_dir,name)) + else: + name_dict[id][1].append(os.path.join(cuhk01_dir,name)) + return name_dict \ No newline at end of file diff --git a/datafolder/reid_dataset/import_CUHK03.py b/datafolder/reid_dataset/import_CUHK03.py new file mode 100644 index 0000000..abb47fc --- /dev/null +++ b/datafolder/reid_dataset/import_CUHK03.py @@ -0,0 +1,46 @@ +import os +from .reiddataset_downloader import * +def import_CUHK03(dataset_dir, detected = False): + + cuhk03_dir = os.path.join(dataset_dir,'CUHK03') + + if not os.path.exists(cuhk03_dir): + Print('Please Download the CUHK03 Dataset') + + if not detected: + cuhk03_dir = os.path.join(cuhk03_dir , 'labeled') + else: + cuhk03_dir = os.path.join(cuhk03_dir , 'detected') + + campair_list = os.listdir(cuhk03_dir) + #campair_list = ['P1','P2','P3'] + name_dict={} + for campair in campair_list: + cam1_list = [] + cam1_list=os.listdir(os.path.join(cuhk03_dir,campair,'cam1')) + cam2_list=os.listdir(os.path.join(cuhk03_dir,campair,'cam2')) + for file in cam1_list: + id = campair[1:]+'-'+file.split('-')[0] + if id not in name_dict: + name_dict[id]=[] + name_dict[id].append([]) + name_dict[id].append([]) + name_dict[id][0].append(os.path.join(cuhk03_dir,campair,'cam1',file)) + for file in cam2_list: + id = campair[1:]+'-'+file.split('-')[0] + if id not in name_dict: + name_dict[id]=[] + name_dict[id].append([]) + name_dict[id].append([]) + name_dict[id][1].append(os.path.join(cuhk03_dir,campair,'cam2',file)) + return name_dict + +def cuhk03_test(data_dir): + CUHK03_dir = os.path.join(data_dir , 'CUHK03') + f = h5py.File(os.path.join(CUHK03_dir,'cuhk-03.mat')) + test = [] + for i in range(20): + test_set = (np.array(f[f['testsets'][0][i]],dtype='int').T).tolist() + test.append(test_set) + + return test \ No newline at end of file diff --git a/datafolder/reid_dataset/import_DukeMTMC.py b/datafolder/reid_dataset/import_DukeMTMC.py new file mode 100644 index 0000000..8fbce23 --- /dev/null +++ b/datafolder/reid_dataset/import_DukeMTMC.py @@ -0,0 +1,32 @@ +import os +from .reiddataset_downloader import* +def import_DukeMTMC(dataset_dir): + dukemtmc_dir = os.path.join(dataset_dir, 'DukeMTMC-reID') + if not os.path.exists(dukemtmc_dir): + print('Please Download the DukMTMC Dataset') + data_group = ['train','query','gallery'] + for group in data_group: + if group == 'train': + name_dir = os.path.join(dukemtmc_dir , 'bounding_box_train') + elif group == 'query': + name_dir = os.path.join(dukemtmc_dir, 'query') + else: + name_dir = os.path.join(dukemtmc_dir, 'bounding_box_test') + file_list=os.listdir(name_dir) + globals()[group]={} + for name in file_list: + if name[-3:]=='jpg': + id = name.split('_')[0] + if id not in globals()[group]: + globals()[group][id]=[] + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + cam_n = int(name.split('_')[1][1])-1 + globals()[group][id][cam_n].append(os.path.join(name_dir,name)) + return train, query, gallery \ No newline at end of file diff --git a/datafolder/reid_dataset/import_DukeMTMCAttribute.py b/datafolder/reid_dataset/import_DukeMTMCAttribute.py new file mode 100644 index 0000000..20902e6 --- /dev/null +++ b/datafolder/reid_dataset/import_DukeMTMCAttribute.py @@ -0,0 +1,126 @@ +import os +from .reiddataset_downloader import * +from .import_DukeMTMC import * +import scipy.io + +def import_DukeMTMCAttribute(dataset_dir): + dataset_name = 'DukeMTMC-reID/attribute' + train,query,test = import_DukeMTMC(dataset_dir) + if not os.path.exists(os.path.join(dataset_dir,dataset_name)): + print('Please Download the DukeMTMCATTributes Dataset') + train_label = ['backpack', + 'bag', + 'handbag', + 'boots', + 'gender', + 'hat', + 'shoes', + 'top', + 'downblack', + 'downwhite', + 'downred', + 'downgray', + 'downblue', + 'downgreen', + 'downbrown', + 'upblack', + 'upwhite', + 'upred', + 'uppurple', + 'upgray', + 'upblue', + 'upgreen', + 'upbrown'] + + test_label=['boots', + 'shoes', + 'top', + 'gender', + 'hat', + 'backpack', + 'bag', + 'handbag', + 'downblack', + 'downwhite', + 'downred', + 'downgray', + 'downblue', + 'downgreen', + 'downbrown', + 'upblack', + 'upwhite', + 'upred', + 'upgray', + 'upblue', + 'upgreen', + 'uppurple', + 'upbrown'] + + + train_person_id = [] + for personid in train: + train_person_id.append(personid) + train_person_id.sort(key=int) + + test_person_id = [] + for personid in test: + test_person_id.append(personid) + test_person_id.sort(key=int) + + f = scipy.io.loadmat(os.path.join(dataset_dir,dataset_name,'duke_attribute.mat')) + + test_attribute = {} + train_attribute = {} + for test_train in range(len(f['duke_attribute'][0][0])): + if test_train == 1: + id_list_name = 'test_person_id' + group_name = 'test_attribute' + else: + id_list_name = 'train_person_id' + group_name = 'train_attribute' + for attribute_id in range(len(f['duke_attribute'][0][0][test_train][0][0])): + if isinstance(f['duke_attribute'][0][0][test_train][0][0][attribute_id][0][0], np.ndarray): + continue + for person_id in range(len(f['duke_attribute'][0][0][test_train][0][0][attribute_id][0])): + id = locals()[id_list_name][person_id] + if id not in locals()[group_name]: + locals()[group_name][id]=[] + locals()[group_name][id].append(f['duke_attribute'][0][0][test_train][0][0][attribute_id][0][person_id]) + + for i in range(8): + train_label.insert(8,train_label[-1]) + train_label.pop(-1) + + unified_train_atr = {} + for k,v in train_attribute.items(): + temp_atr = list(v) + for i in range(8): + temp_atr.insert(8,temp_atr[-1]) + temp_atr.pop(-1) + unified_train_atr[k] = temp_atr + + unified_test_atr = {} + for k,v in test_attribute.items(): + temp_atr = [0]*len(train_label) + for i in range(len(train_label)): + temp_atr[i]=v[test_label.index(train_label[i])] + unified_test_atr[k] = temp_atr + #two zero appear in train '0370' '0679' + #zero_check=[] + #for id in train_attribute: + # if 0 in train_attribute[id]: + # zero_check.append(id) + #for i in range(len(zero_check)): + # train_attribute[zero_check[i]] = [1 if x==0 else x for x in train_attribute[zero_check[i]]] + unified_train_atr['0370'][7]=1 + unified_train_atr['0679'][7]=2 + + return unified_train_atr,unified_test_atr,train_label + +def import_DukeMTMCAttribute_binary(dataset_dir): + train_duke_attr, test_duke_attr,label = import_DukeMTMCAttribute(dataset_dir) + for id in train_duke_attr: + train_duke_attr[id][:] = [x - 1 for x in train_duke_attr[id]] + for id in test_duke_attr: + test_duke_attr[id][:] = [x - 1 for x in test_duke_attr[id]] + return train_duke_attr, test_duke_attr, label diff --git a/datafolder/reid_dataset/import_Market1501.py b/datafolder/reid_dataset/import_Market1501.py new file mode 100644 index 0000000..8f315a8 --- /dev/null +++ b/datafolder/reid_dataset/import_Market1501.py @@ -0,0 +1,30 @@ +import os +from .reiddataset_downloader import * +def import_Market1501(dataset_dir): + market1501_dir = os.path.join(dataset_dir,'Market-1501') + if not os.path.exists(market1501_dir): + print('Please Download Market1501 Dataset') + data_group = ['train','query','gallery'] + for group in data_group: + if group == 'train': + name_dir = os.path.join(market1501_dir , 'bounding_box_train') + elif group == 'query': + name_dir = os.path.join(market1501_dir, 'query') + else: + name_dir = os.path.join(market1501_dir, 'bounding_box_test') + file_list=os.listdir(name_dir) + globals()[group]={} + for name in file_list: + if name[-3:]=='jpg': + id = name.split('_')[0] + if id not in globals()[group]: + globals()[group][id]=[] + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + globals()[group][id].append([]) + cam_n = int(name.split('_')[1][1])-1 + globals()[group][id][cam_n].append(os.path.join(name_dir,name)) + return train,query,gallery \ No newline at end of file diff --git a/datafolder/reid_dataset/import_Market1501Attribute.py b/datafolder/reid_dataset/import_Market1501Attribute.py new file mode 100644 index 0000000..2b44d01 --- /dev/null +++ b/datafolder/reid_dataset/import_Market1501Attribute.py @@ -0,0 +1,174 @@ +import os +from .import_Market1501 import * +from .reiddataset_downloader import * +import scipy.io + + +def import_Market1501Attribute(dataset_dir): + dataset_name = 'Market-1501/attribute' + train,query,test = import_Market1501(dataset_dir) + if not os.path.exists(os.path.join(dataset_dir,dataset_name)): + print('Please Download the Market1501Attribute Dataset') + train_label=['age', + 'backpack', + 'bag', + 'handbag', + 'downblack', + 'downblue', + 'downbrown', + 'downgray', + 'downgreen', + 'downpink', + 'downpurple', + 'downwhite', + 'downyellow', + 'upblack', + 'upblue', + 'upgreen', + 'upgray', + 'uppurple', + 'upred', + 'upwhite', + 'upyellow', + 'clothes', + 'down', + 'up', + 'hair', + 'hat', + 'gender'] + + test_label=['age', + 'backpack', + 'bag', + 'handbag', + 'clothes', + 'down', + 'up', + 'hair', + 'hat', + 'gender', + 'upblack', + 'upwhite', + 'upred', + 'uppurple', + 'upyellow', + 'upgray', + 'upblue', + 'upgreen', + 'downblack', + 'downwhite', + 'downpink', + 'downpurple', + 'downyellow', + 'downgray', + 'downblue', + 'downgreen', + 'downbrown' + ] + + train_person_id = [] + for personid in train: + train_person_id.append(personid) + train_person_id.sort(key=int) + + test_person_id = [] + for personid in test: + test_person_id.append(personid) + test_person_id.sort(key=int) + test_person_id.remove('-1') + test_person_id.remove('0000') + + f = scipy.io.loadmat(os.path.join(dataset_dir,dataset_name,'market_attribute.mat')) + + test_attribute = {} + train_attribute = {} + for test_train in range(len(f['market_attribute'][0][0])): + if test_train == 0: + id_list_name = 'test_person_id' + group_name = 'test_attribute' + else: + id_list_name = 'train_person_id' + group_name = 'train_attribute' + for attribute_id in range(len(f['market_attribute'][0][0][test_train][0][0])): + if isinstance(f['market_attribute'][0][0][test_train][0][0][attribute_id][0][0], np.ndarray): + continue + for person_id in range(len(f['market_attribute'][0][0][test_train][0][0][attribute_id][0])): + id = locals()[id_list_name][person_id] + if id not in locals()[group_name]: + locals()[group_name][id]=[] + locals()[group_name][id].append(f['market_attribute'][0][0][test_train][0][0][attribute_id][0][person_id]) + + unified_train_atr = {} + for k,v in train_attribute.items(): + temp_atr = [0]*len(test_label) + for i in range(len(test_label)): + temp_atr[i]=v[train_label.index(test_label[i])] + unified_train_atr[k] = temp_atr + + return unified_train_atr, test_attribute, test_label + + +def import_Market1501Attribute_binary(dataset_dir): + train_market_attr, test_market_attr, label = import_Market1501Attribute(dataset_dir) + + for id in train_market_attr: + train_market_attr[id][:] = [x - 1 for x in train_market_attr[id]] + if train_market_attr[id][0] == 0: + train_market_attr[id].pop(0) + train_market_attr[id].insert(0, 1) + train_market_attr[id].insert(1, 0) + train_market_attr[id].insert(2, 0) + train_market_attr[id].insert(3, 0) + elif train_market_attr[id][0] == 1: + train_market_attr[id].pop(0) + train_market_attr[id].insert(0, 0) + train_market_attr[id].insert(1, 1) + train_market_attr[id].insert(2, 0) + train_market_attr[id].insert(3, 0) + elif train_market_attr[id][0] == 2: + train_market_attr[id].pop(0) + train_market_attr[id].insert(0, 0) + train_market_attr[id].insert(1, 0) + train_market_attr[id].insert(2, 1) + train_market_attr[id].insert(3, 0) + elif train_market_attr[id][0] == 3: + train_market_attr[id].pop(0) + train_market_attr[id].insert(0, 0) + train_market_attr[id].insert(1, 0) + train_market_attr[id].insert(2, 0) + train_market_attr[id].insert(3, 1) + + for id in test_market_attr: + test_market_attr[id][:] = [x - 1 for x in test_market_attr[id]] + if test_market_attr[id][0] == 0: + test_market_attr[id].pop(0) + test_market_attr[id].insert(0, 1) + test_market_attr[id].insert(1, 0) + test_market_attr[id].insert(2, 0) + test_market_attr[id].insert(3, 0) + elif test_market_attr[id][0] == 1: + test_market_attr[id].pop(0) + test_market_attr[id].insert(0, 0) + test_market_attr[id].insert(1, 1) + test_market_attr[id].insert(2, 0) + test_market_attr[id].insert(3, 0) + elif test_market_attr[id][0] == 2: + test_market_attr[id].pop(0) + test_market_attr[id].insert(0, 0) + test_market_attr[id].insert(1, 0) + test_market_attr[id].insert(2, 1) + test_market_attr[id].insert(3, 0) + elif test_market_attr[id][0] == 3: + test_market_attr[id].pop(0) + test_market_attr[id].insert(0, 0) + test_market_attr[id].insert(1, 0) + test_market_attr[id].insert(2, 0) + test_market_attr[id].insert(3, 1) + + label.pop(0) + label.insert(0,'young') + label.insert(1,'teenager') + label.insert(2,'adult') + label.insert(3,'old') + + return train_market_attr, test_market_attr, label \ No newline at end of file diff --git a/datafolder/reid_dataset/import_MarketDuke.py b/datafolder/reid_dataset/import_MarketDuke.py new file mode 100644 index 0000000..11c961e --- /dev/null +++ b/datafolder/reid_dataset/import_MarketDuke.py @@ -0,0 +1,32 @@ +import os +from .reiddataset_downloader import * + + +def import_MarketDuke(data_dir, dataset_name): + dataset_dir = os.path.join(data_dir,dataset_name) + + if not os.path.exists(dataset_dir): + print('Please Download '+dataset_name+ ' Dataset') + + dataset_dir = os.path.join(data_dir,dataset_name) + data_group = ['train','query','gallery'] + for group in data_group: + if group == 'train': + name_dir = os.path.join(dataset_dir , 'bounding_box_train') + elif group == 'query': + name_dir = os.path.join(dataset_dir, 'query') + else: + name_dir = os.path.join(dataset_dir, 'bounding_box_test') + file_list=sorted(os.listdir(name_dir)) + globals()[group]={} + globals()[group]['data']=[] + globals()[group]['ids'] = [] + for name in file_list: + if name[-3:]=='jpg': + id = name.split('_')[0] + cam = int(name.split('_')[1][1]) + images = os.path.join(name_dir,name) + if id not in globals()[group]['ids']: + globals()[group]['ids'].append(id) + globals()[group]['data'].append([images,globals()[group]['ids'].index(id),id,cam,name.split('.')[0]]) + return train,query,gallery \ No newline at end of file diff --git a/datafolder/reid_dataset/import_MarketDuke_nodistractors.py b/datafolder/reid_dataset/import_MarketDuke_nodistractors.py new file mode 100644 index 0000000..d3f4dad --- /dev/null +++ b/datafolder/reid_dataset/import_MarketDuke_nodistractors.py @@ -0,0 +1,33 @@ +import os +from .reiddataset_downloader import * + + +def import_MarketDuke_nodistractors(data_dir, dataset_name): + dataset_dir = os.path.join(data_dir,dataset_name) + + if not os.path.exists(dataset_dir): + print('Please Download '+dataset_name+ ' Dataset') + + dataset_dir = os.path.join(data_dir,dataset_name) + data_group = ['train','query','gallery'] + for group in data_group: + if group == 'train': + name_dir = os.path.join(dataset_dir , 'bounding_box_train') + elif group == 'query': + name_dir = os.path.join(dataset_dir, 'query') + else: + name_dir = os.path.join(dataset_dir, 'bounding_box_test') + file_list=sorted(os.listdir(name_dir)) + globals()[group]={} + globals()[group]['data']=[] + globals()[group]['ids'] = [] + for name in file_list: + if name[-3:]=='jpg': + id = name.split('_')[0] + cam = int(name.split('_')[1][1]) + images = os.path.join(name_dir,name) + if (id!='0000' and id !='-1'): + if id not in globals()[group]['ids']: + globals()[group]['ids'].append(id) + globals()[group]['data'].append([images,globals()[group]['ids'].index(id),id,cam,name.split('.')[0]]) + return train, query, gallery \ No newline at end of file diff --git a/datafolder/reid_dataset/import_VIPeR.py b/datafolder/reid_dataset/import_VIPeR.py new file mode 100644 index 0000000..fa22462 --- /dev/null +++ b/datafolder/reid_dataset/import_VIPeR.py @@ -0,0 +1,29 @@ +import os +from .reiddataset_downloader import * +def import_VIPeR(dataset_dir): + viper_dir = os.path.join(dataset_dir , 'VIPeR') + if not os.path.exists(viper_dir): + orint('Please Download VIPeR Dataset') + + file_list_a=os.listdir(os.path.join(viper_dir,'cam_a')) + file_list_b=os.listdir(os.path.join(viper_dir,'cam_b')) + + name_dict={} + for name in file_list_a: + if name[-3:]=='bmp': + id = name.split('_')[0] + if id not in name_dict: + name_dict[id]=[] + name_dict[id].append([]) + name_dict[id].append([]) + name_dict[id][0].append(os.path.join(viper_dir,'cam_a',name)) + for name in file_list_b: + if name[-3:]=='bmp': + id = name.split('_')[0] + if id not in name_dict: + name_dict[id]=[] + name_dict[id].append([]) + name_dict[id].append([]) + name_dict[id][1].append(os.path.join(viper_dir,'cam_b',name)) + + return name_dict \ No newline at end of file diff --git a/datafolder/reid_dataset/marketduke_to_hdf5.py b/datafolder/reid_dataset/marketduke_to_hdf5.py new file mode 100644 index 0000000..3c229f6 --- /dev/null +++ b/datafolder/reid_dataset/marketduke_to_hdf5.py @@ -0,0 +1,32 @@ +import warnings +warnings.filterwarnings('ignore','.*conversion.*') + +import os +import h5py +import numpy as np +from PIL import Image +from .import_MarketDuke import import_MarketDuke + +def marketduke_to_hdf5(data_dir,dataset_name,save_dir=os.getcwd()): + phase_list = ['train','query','gallery'] + dataset = import_MarketDuke(data_dir,dataset_name) + dt = h5py.special_dtype(vlen=str) + + f = h5py.File(os.path.join(save_dir,dataset_name+'.hdf5'),'w') + for phase in phase_list: + grp = f.create_group(phase) + phase_dataset = dataset[phase_list.index(phase)] + for i in range(len(phase_dataset['data'])): + name = phase_dataset['data'][i][0].split('/')[-1].split('.')[0] + temp = grp.create_group(name) + temp.create_dataset('img',data=Image.open(phase_dataset['data'][i][0])) + temp.create_dataset('index',data=int(phase_dataset['data'][i][1])) + temp.create_dataset('id',data=phase_dataset['data'][i][2], dtype=dt) + temp.create_dataset('cam',data=int(phase_dataset['data'][i][3])) + + ids = f.create_group('ids') + ids.create_dataset('train',data=np.array(dataset[0]['ids'],'S4'),dtype=dt) + ids.create_dataset('query',data=np.array(dataset[1]['ids'],'S4'),dtype=dt) + ids.create_dataset('gallery',data=np.array(dataset[2]['ids'],'S4'),dtype=dt) + + f.close() \ No newline at end of file diff --git a/datafolder/reid_dataset/pytorch_prepare.py b/datafolder/reid_dataset/pytorch_prepare.py new file mode 100644 index 0000000..1f09495 --- /dev/null +++ b/datafolder/reid_dataset/pytorch_prepare.py @@ -0,0 +1,96 @@ +import os +from shutil import copyfile + +def pytorch_prepare(data_dir, dataset_name): + dataset_dir = os.path.join(data_dir, dataset_name) + + if not os.path.isdir(dataset_dir): + print('please change the download_path') + + pytorch_path = os.path.join(dataset_dir , 'pytorch') + + if not os.path.isdir(pytorch_path): + os.mkdir(pytorch_path) + #----------------------------------------- + #query + print('generatring ' + dataset_name + ' query images.') + query_dir = os.path.join(dataset_dir , 'query') + query_save_dir = os.path.join(dataset_dir , 'pytorch', 'query') + if not os.path.isdir(query_save_dir): + os.mkdir(query_save_dir) + + for root, dirs, files in os.walk(query_dir, topdown=True): + for name in files: + if not name[-3:]=='jpg': + continue + ID = name.split('_') + src_dir = os.path.join(query_dir , name) + dst_dir = os.path.join(query_save_dir, ID[0]) + if not os.path.isdir(dst_dir): + os.mkdir(dst_dir) + copyfile(src_dir, os.path.join(dst_dir , name)) + #----------------------------------------- + #gallery + print('generatring '+dataset_name+' gallery images.') + gallery_dir = os.path.join(dataset_dir , 'bounding_box_test') + gallery_save_dir = os.path.join(dataset_dir , 'pytorch' , 'gallery') + if not os.path.isdir(gallery_save_dir): + os.mkdir(gallery_save_dir) + + for root, dirs, files in os.walk(gallery_dir, topdown=True): + for name in files: + if not name[-3:]=='jpg': + continue + ID = name.split('_') + src_dir = os.path.join(gallery_dir, name) + dst_dir = os.path.join(gallery_save_dir, ID[0]) + if not os.path.isdir(dst_dir): + os.mkdir(dst_dir) + copyfile(src_dir, os.path.join(dst_dir,name)) + #--------------------------------------- + #train_all + print('generatring '+dataset_name + ' all training images.') + train_dir = os.path.join( dataset_dir , 'bounding_box_train') + train_save_all_dir = os.path.join( dataset_dir , 'pytorch', 'train_all') + if not os.path.isdir(train_save_all_dir): + os.mkdir(train_save_all_dir) + + for root, dirs, files in os.walk(train_dir, topdown=True): + for name in files: + if not name[-3:]=='jpg': + continue + ID = name.split('_') + src_dir = os.path.join(train_dir , name) + dst_dir = os.path.join(train_save_all_dir, ID[0]) + if not os.path.isdir(dst_dir): + os.mkdir(dst_dir) + copyfile(src_dir, os.path.join(dst_dir, name)) + + #--------------------------------------- + #train_val + print('generatring '+ dataset_name+' training and validation images.') + train_save_dir = os.path.join(dataset_dir, 'pytorch', 'train') + val_save_dir = os.path.join(dataset_dir , 'pytorch' , 'val') + if not os.path.isdir(train_save_dir): + os.mkdir(train_save_dir) + os.mkdir(val_save_dir) + + for root, dirs, files in os.walk(train_dir, topdown=True): + for name in files: + if not name[-3:]=='jpg': + continue + ID = name.split('_') + src_dir = os.path.join(train_dir , name) + dst_dir = os.path.join(train_save_dir , ID[0]) + if not os.path.isdir(dst_dir): + os.mkdir(dst_dir) + dst_dir = os.path.join(val_save_dir, ID[0]) #first image is used as val image + os.mkdir(dst_dir) + copyfile(src_dir, os.path.join(dst_dir , name)) + print('Finished ' + dataset_name) + else: + print(dataset_name + ' pytorch directory exists!') + +def pytorch_prepare_all(data_dir): + pytorch_prepare('Market1501', data_dir) + pytorch_prepare('DukeMTMC', data_dir) diff --git a/datafolder/reid_dataset/reiddataset_downloader.py b/datafolder/reid_dataset/reiddataset_downloader.py new file mode 100644 index 0000000..1e84fec --- /dev/null +++ b/datafolder/reid_dataset/reiddataset_downloader.py @@ -0,0 +1,88 @@ +from __future__ import print_function +import warnings +warnings.filterwarnings('ignore','.*conversion.*') + +import os +import zipfile +import shutil +import requests +import h5py +import numpy as np +from PIL import Image +import argparse +from .gdrive_downloader import gdrive_downloader +from .cuhk03_to_image import cuhk03_to_image + +dataset = { + 'CUHK01': '153IzD3vyQ0PqxxanQRlP9l89F1S5Vr47', + 'CUHK02': '0B2FnquNgAXoneE5YamFXY3NjYWM', + 'CUHK03': '1BO4G9gbOTJgtYIB0VNyHQpZb8Lcn-05m', + 'VIPeR': '0B2FnquNgAXonZzJPQUtrcWJWbWc', + 'Market1501': '0B2FnquNgAXonU3RTcE1jQlZ3X0E', + 'Market1501Attribute' : '1YMgni5oz-RPkyKHzOKnYRR2H3IRKdsHO', + 'DukeMTMC': '1qtFGJQ6eFu66Tt7WG85KBxtACSE8RBZ0', + 'DukeMTMCAttribute' : '1eilPJFnk_EHECKj2glU_ZLLO7eR3JIiO' +} + +dataset_hdf5 = { + 'Market1501': '1ipvyt4qesVK6CUiGcQdwle2c2XYknKco', + 'DukeMTMC': '1tP-fty5YE-W2F6B5rjnQNfE-NzNssGM2' +} + +def reiddataset_downloader(data_dir, data_name, hdf5 = True): + + if not os.path.exists(data_dir): + os.makedirs(data_dir) + + if hdf5: + dataset_dir = os.path.join(data_dir , data_name) + if not os.path.exists(dataset_dir): + os.makedirs(dataset_dir) + destination = os.path.join(dataset_dir , data_name+'.hdf5') + if not os.path.isfile(destination): + id = dataset_hdf5[data_name] + print("Downloading %s in HDF5 Formate" %data_name) + gdrive_downloader(destination, id) + print("Done") + else: + print("Dataset Check Success: %s exists!" %data_name) + else: + data_dir_exist = os.path.join(data_dir , data_name) + + if not os.path.exists(data_dir_exist): + temp_dir = os.path.join(data_dir , 'temp') + + if not os.path.exists(temp_dir): + os.makedirs(temp_dir) + + destination = os.path.join(temp_dir , data_name) + + id = dataset[data_name] + + print("Downloading %s in Original Images" % data_name) + gdrive_downloader(destination, id) + + zip_ref = zipfile.ZipFile(destination) + print("Extracting %s" % data_name) + zip_ref.extractall(data_dir) + zip_ref.close() + shutil.rmtree(temp_dir) + print("Done") + if data_name == 'CUHK03': + print('Converting cuhk03.mat into images') + cuhk03_to_image(os.path.join(data_dir,'CUHK03')) + print('Done') + else: + print("Dataset Check Success: %s exists!" %data_name) + +def reiddataset_downloader_all(data_dir): + for k,v in dataset.items(): + reiddataset_downloader(k,data_dir) + +#For United Testing and External Use +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Dataset Name and Dataset Directory') + parser.add_argument(dest="data_dir", action="store", default="~/Datasets/",help="") + parser.add_argument(dest="data_name", action="store", type=str,help="") + args = parser.parse_args() + reiddataset_downloader(args.data_dir,args.data_name) \ No newline at end of file diff --git a/doc/DukeMTMC-reID-attribute b/doc/DukeMTMC-reID-attribute new file mode 100644 index 0000000..16289e9 --- /dev/null +++ b/doc/DukeMTMC-reID-attribute @@ -0,0 +1,16 @@ +The 23 attributes are: + +| attribute | representation in file | label | + +| gender | gender | male(1), female(2) | +| length of upper-body clothing | top | short upper body clothing(1), long(2) | +| wearing boots | boots | no(1), yes(2) | +| wearing hat | hat | no(1), yes(2) | +| carrying backpack | backpack | no(1), yes(2) | +| carrying bag | bag | no(1), yes(2) | +| carrying handbag | handbag | no(1), yes(2) | +| color of shoes | shoes | dark(1), light(2) | +| 8 color of upper-body clothing | upblack, upwhite, upred, uppurple, upgray, upblue, upgreen, upbrown | no(1), yes(2) | +| 7 color of lower-body clothing | downblack, downwhite, downred, downgray, downblue, downgreen, downbrown | no(1), yes(2) | + +Note that the though there are 7 and 8 attributes for lower-body clothing and upper-body clothing, only one color is labeled as yes (2) for an identity. diff --git a/doc/Market1501-attribute b/doc/Market1501-attribute new file mode 100644 index 0000000..8c7c705 --- /dev/null +++ b/doc/Market1501-attribute @@ -0,0 +1,21 @@ +The 27 attributes are: + +| attribute | representation in file | label | +| :----------------------------: | :--------------------: | :-----------------------------------------: | +| gender | gender | male(1), female(2) | +| hair length | hair | short hair(1), long hair(2) | +| sleeve length | up | long sleeve(1), short sleeve(2) | +| length of lower-body clothing | down | long lower body clothing(1), short(2) | +| type of lower-body clothing | clothes | dress(1), pants(2) | +| wearing hat | hat | no(1), yes(2) | +| carrying backpack | backpack | no(1), yes(2) | +| carrying bag | bag | no(1), yes(2) | +| carrying handbag | handbag | no(1), yes(2) | +| age | age | young(1), teenager(2), adult(3), old(4) | + + +| 8 color of upper-body clothing | upblack, upwhite, upred, uppurple, upyellow, upgray, upblue, upgreen | no(1), yes(2) | +| 9 color of lower-body clothing | downblack, downwhite, downpink, downpurple, downyellow, downgray, downblue, downgreen,downbrown | no(1), yes(2) | + +Note that the though there are 8 and 9 attributes for upper-body clothing and lower-body clothing, only one color is labeled as yes (2) for an identity. + diff --git a/doc/attribute.json b/doc/attribute.json new file mode 100644 index 0000000..fb4858c --- /dev/null +++ b/doc/attribute.json @@ -0,0 +1,377 @@ +{ + "duke": { + "bag": [ + "carrying bag", + [ + "no", + "yes" + ] + ], + "upred": [ + "color of upper-body clothing", + [ + null, + "red" + ] + ], + "upblue": [ + "color of upper-body clothing", + [ + null, + "blue" + ] + ], + "hat": [ + "wearing hat", + [ + "no", + "yes" + ] + ], + "shoes": [ + "color of shoes", + [ + "dark", + "light" + ] + ], + "downgreen": [ + "color of lower-body clothing", + [ + null, + "green" + ] + ], + "downbrown": [ + "color of lower-body clothing", + [ + null, + "brown" + ] + ], + "downred": [ + "color of lower-body clothing", + [ + null, + "red" + ] + ], + "upgreen": [ + "color of upper-body clothing", + [ + null, + "green" + ] + ], + "handbag": [ + "carrying handbag", + [ + "no", + "yes" + ] + ], + "downblue": [ + "color of lower-body clothing", + [ + null, + "blue" + ] + ], + "downblack": [ + "color of lower-body clothing", + [ + null, + "black" + ] + ], + "backpack": [ + "carrying backpack", + [ + "no", + "yes" + ] + ], + "downwhite": [ + "color of lower-body clothing", + [ + null, + "white" + ] + ], + "upblack": [ + "color of upper-body clothing", + [ + null, + "black" + ] + ], + "gender": [ + "gender", + [ + "male", + "female" + ] + ], + "downgray": [ + "color of lower-body clothing", + [ + null, + "gray" + ] + ], + "boots": [ + "wearing boots", + [ + "no", + "yes" + ] + ], + "uppurple": [ + "color of upper-body clothing", + [ + null, + "purple" + ] + ], + "top": [ + "length of upper-body clothing", + [ + "short upper body clothing", + "long" + ] + ], + "upbrown": [ + "color of upper-body clothing", + [ + null, + "brown" + ] + ], + "upgray": [ + "color of upper-body clothing", + [ + null, + "gray" + ] + ], + "upwhite": [ + "color of upper-body clothing", + [ + null, + "white" + ] + ] + }, + "market": { + "bag": [ + "carrying bag", + [ + "no", + "yes" + ] + ], + "upred": [ + "color of upper-body clothing", + [ + null, + "red" + ] + ], + "upblue": [ + "color of upper-body clothing", + [ + null, + "blue" + ] + ], + "hat": [ + "wearing hat", + [ + "no", + "yes" + ] + ], + "downgreen": [ + "color of lower-body clothing", + [ + null, + "green" + ] + ], + "downbrown": [ + "color of lower-body clothing", + [ + null, + "brown" + ] + ], + "upyellow": [ + "color of upper-body clothing", + [ + null, + "yellow" + ] + ], + "up": [ + "sleeve length", + [ + "long sleeve", + "short sleeve" + ] + ], + "upgreen": [ + "color of upper-body clothing", + [ + null, + "green" + ] + ], + "handbag": [ + "carrying handbag", + [ + "no", + "yes" + ] + ], + "downgray": [ + "color of lower-body clothing", + [ + null, + "gray" + ] + ], + "clothes": [ + "type of lower-body clothing", + [ + "dress", + "pants" + ] + ], + "adult": [ + "age", + [ + null, + "adult" + ] + ], + "downblack": [ + "color of lower-body clothing", + [ + null, + "black" + ] + ], + "backpack": [ + "carrying backpack", + [ + "no", + "yes" + ] + ], + "downwhite": [ + "color of lower-body clothing", + [ + null, + "white" + ] + ], + "upblack": [ + "color of upper-body clothing", + [ + null, + "black" + ] + ], + "gender": [ + "gender", + [ + "male", + "female" + ] + ], + "downyellow": [ + "color of lower-body clothing", + [ + null, + "yellow" + ] + ], + "downpink": [ + "color of lower-body clothing", + [ + null, + "pink" + ] + ], + "old": [ + "age", + [ + null, + "old" + ] + ], + "down": [ + "length of lower-body clothing", + [ + "long lower body clothing", + "short" + ] + ], + "uppurple": [ + "color of upper-body clothing", + [ + null, + "purple" + ] + ], + "downpurple": [ + "color of lower-body clothing", + [ + null, + "purple" + ] + ], + "young": [ + "age", + [ + null, + "young" + ] + ], + "teenager": [ + "age", + [ + null, + "teenager" + ] + ], + "hair": [ + "hair length", + [ + "short hair", + "long hair" + ] + ], + "downblue": [ + "color of lower-body clothing", + [ + null, + "blue" + ] + ], + "upgray": [ + "color of upper-body clothing", + [ + null, + "gray" + ] + ], + "upwhite": [ + "color of upper-body clothing", + [ + null, + "white" + ] + ] + } +} \ No newline at end of file diff --git a/doc/label.json b/doc/label.json new file mode 100644 index 0000000..fa81dcb --- /dev/null +++ b/doc/label.json @@ -0,0 +1,59 @@ +{ + "duke": [ + "backpack", + "bag", + "handbag", + "boots", + "gender", + "hat", + "shoes", + "top", + "upblack", + "upwhite", + "upred", + "uppurple", + "upgray", + "upblue", + "upgreen", + "upbrown", + "downblack", + "downwhite", + "downred", + "downgray", + "downblue", + "downgreen", + "downbrown" + ], + "market": [ + "young", + "teenager", + "adult", + "old", + "backpack", + "bag", + "handbag", + "clothes", + "down", + "up", + "hair", + "hat", + "gender", + "upblack", + "upwhite", + "upred", + "uppurple", + "upyellow", + "upgray", + "upblue", + "upgreen", + "downblack", + "downwhite", + "downpink", + "downpurple", + "downyellow", + "downgray", + "downblue", + "downgreen", + "downbrown" + ] +} \ No newline at end of file diff --git a/doc/save_json.py b/doc/save_json.py new file mode 100644 index 0000000..cab5e4b --- /dev/null +++ b/doc/save_json.py @@ -0,0 +1,141 @@ +import json + + +market_label_list = [ + "young", + "teenager", + "adult", + "old", + "backpack", + "bag", + "handbag", + "clothes", + "down", + "up", + "hair", + "hat", + "gender", + "upblack", + "upwhite", + "upred", + "uppurple", + "upyellow", + "upgray", + "upblue", + "upgreen", + "downblack", + "downwhite", + "downpink", + "downpurple", + "downyellow", + "downgray", + "downblue", + "downgreen", + "downbrown" +] + + +duke_label_list = [ + "backpack", + "bag", + "handbag", + "boots", + "gender", + "hat", + "shoes", + "top", + "upblack", + "upwhite", + "upred", + "uppurple", + "upgray", + "upblue", + "upgreen", + "upbrown", + "downblack", + "downwhite", + "downred", + "downgray", + "downblue", + "downgreen", + "downbrown" +] + + +market_attribute_dict = { + "young": ["age", [None, "young"]], + "teenager": ["age", [None, "teenager"]], + "adult": ["age", [None, "adult"]], + "old": ["age", [None, "old"]], + "backpack": ["carrying backpack", ["no", "yes"]], + "bag": ["carrying bag", ["no", "yes"]], + "handbag": ["carrying handbag", ["no", "yes"]], + "clothes": ["type of lower-body clothing", ["dress", "pants"]], + "down": ["length of lower-body clothing", ["long lower body clothing", "short"]], + "up": ["sleeve length", ["long sleeve", "short sleeve"]], + "hair": ["hair length", ["short hair", "long hair"]], + "hat": ["wearing hat", ["no", "yes"]], + "gender": ["gender", ["male", "female"]], + "upblack": ["color of upper-body clothing", [None, "black"]], + "upwhite": ["color of upper-body clothing", [None, "white"]], + "upred": ["color of upper-body clothing", [None, "red"]], + "uppurple": ["color of upper-body clothing", [None, "purple"]], + "upyellow": ["color of upper-body clothing", [None, "yellow"]], + "upgray": ["color of upper-body clothing", [None, "gray"]], + "upblue": ["color of upper-body clothing", [None, "blue"]], + "upgreen": ["color of upper-body clothing", [None, "green"]], + "downblack": ["color of lower-body clothing", [None, "black"]], + "downwhite": ["color of lower-body clothing", [None, "white"]], + "downpink": ["color of lower-body clothing", [None, "pink"]], + "downpurple": ["color of lower-body clothing", [None, "purple"]], + "downyellow": ["color of lower-body clothing", [None, "yellow"]], + "downgray": ["color of lower-body clothing", [None, "gray"]], + "downblue": ["color of lower-body clothing", [None, "blue"]], + "downgreen": ["color of lower-body clothing", [None, "green"]], + "downbrown": ["color of lower-body clothing", [None, "brown"]], +} + + +duke_attribute_dict = { + "backpack": ["carrying backpack", ["no", "yes"]], + "bag": ["carrying bag", ["no", "yes"]], + "handbag": ["carrying handbag", ["no", "yes"]], + "boots": ["wearing boots", ["no", "yes"]], + "gender": ["gender", ["male", "female"]], + "hat": ["wearing hat", ["no", "yes"]], + "shoes": ["color of shoes", ["dark", "light"]], + "top": ["length of upper-body clothing", ["short upper body clothing", "long"]], + "upblack": ["color of upper-body clothing", [None, "black"]], + "upwhite": ["color of upper-body clothing", [None, "white"]], + "upred": ["color of upper-body clothing", [None, "red"]], + "uppurple": ["color of upper-body clothing", [None, "purple"]], + "upgray": ["color of upper-body clothing", [None, "gray"]], + "upblue": ["color of upper-body clothing", [None, "blue"]], + "upgreen": ["color of upper-body clothing", [None, "green"]], + "upbrown": ["color of upper-body clothing", [None, "brown"]], + "downblack": ["color of lower-body clothing", [None, "black"]], + "downwhite": ["color of lower-body clothing", [None, "white"]], + "downred": ["color of lower-body clothing", [None, "red"]], + "downgray": ["color of lower-body clothing", [None, "gray"]], + "downblue": ["color of lower-body clothing", [None, "blue"]], + "downgreen": ["color of lower-body clothing", [None, "green"]], + "downbrown": ["color of lower-body clothing", [None, "brown"]], +} + + +with open('./label.json', 'w') as f: + label_list_dict = { + 'market': market_label_list, + 'duke': duke_label_list, + } + jsObj = json.dumps(label_list_dict, indent=4) + f.write(jsObj) + + +with open('./attribute.json', 'w') as f: + attribute_list_dict = { + 'market': market_attribute_dict, + 'duke': duke_attribute_dict, + } + jsObj = json.dumps(attribute_list_dict, indent=4) + f.write(jsObj) \ No newline at end of file diff --git a/net/__init__.py b/net/__init__.py new file mode 100644 index 0000000..7322f24 --- /dev/null +++ b/net/__init__.py @@ -0,0 +1,9 @@ +from .models import Backbone_nFC, Backbone_nFC_Id + + +def get_model(model_name, num_label, use_id=False, num_id=None): + if not use_id: + return Backbone_nFC(num_label, model_name) + else: + return Backbone_nFC_Id(num_label, num_id, model_name) + diff --git a/net/models.py b/net/models.py new file mode 100644 index 0000000..025b15c --- /dev/null +++ b/net/models.py @@ -0,0 +1,76 @@ +import torch +from torch import nn +from torch.nn import init +from torchvision import models +from net.utils import ClassBlock +from torch.nn import functional as F + + +class Backbone_nFC(nn.Module): + def __init__(self, class_num, model_name='resnet50_nfc'): + super(Backbone_nFC, self).__init__() + self.model_name = model_name + self.backbone_name = model_name.split('_')[0] + self.class_num = class_num + + model_ft = getattr(models, self.backbone_name)(pretrained=True) + if 'resnet' in self.backbone_name: + model_ft.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + model_ft.fc = nn.Sequential() + self.features = model_ft + self.num_ftrs = 2048 + elif 'densenet' in self.backbone_name: + model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + model_ft.fc = nn.Sequential() + self.features = model_ft.features + self.num_ftrs = 1024 + else: + raise NotImplementedError + + for c in range(self.class_num): + self.__setattr__('class_%d' % c, ClassBlock(input_dim=self.num_ftrs, class_num=1, activ='sigmoid') ) + + def forward(self, x): + x = self.features(x) + x = x.view(x.size(0), -1) + pred_label = [self.__getattr__('class_%d' % c)(x) for c in range(self.class_num)] + pred_label = torch.cat(pred_label, dim=1) + return pred_label + + +class Backbone_nFC_Id(nn.Module): + def __init__(self, class_num, id_num, model_name='resnet50_nfc_id'): + super(Backbone_nFC_Id, self).__init__() + self.model_name = model_name + self.backbone_name = model_name.split('_')[0] + self.class_num = class_num + self.id_num = id_num + + model_ft = getattr(models, self.backbone_name)(pretrained=True) + if 'resnet' in self.backbone_name: + model_ft.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + model_ft.fc = nn.Sequential() + self.features = model_ft + self.num_ftrs = 2048 + elif 'densenet' in self.backbone_name: + model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + model_ft.fc = nn.Sequential() + self.features = model_ft.features + self.num_ftrs = 1024 + else: + raise NotImplementedError + + for c in range(self.class_num+1): + if c == self.class_num: + self.__setattr__('class_%d' % c, ClassBlock(self.num_ftrs, class_num=self.id_num, activ='none')) + else: + self.__setattr__('class_%d' % c, ClassBlock(self.num_ftrs, class_num=1, activ='sigmoid')) + + def forward(self, x): + x = self.features(x) + x = x.view(x.size(0), -1) + pred_label = [self.__getattr__('class_%d' % c)(x) for c in range(self.class_num)] + pred_label = torch.cat(pred_label, dim=1) + pred_id = self.__getattr__('class_%d' % self.class_num)(x) + return pred_label, pred_id + diff --git a/net/utils.py b/net/utils.py new file mode 100644 index 0000000..3239e8b --- /dev/null +++ b/net/utils.py @@ -0,0 +1,60 @@ +from torch import nn +from torch.nn import init + + +def weights_init_kaiming(m): + classname = m.__class__.__name__ + # print(classname) + if classname.find('Conv') != -1: + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + elif classname.find('Linear') != -1: + init.kaiming_normal_(m.weight.data, a=0, mode='fan_out') + init.constant_(m.bias.data, 0.0) + elif classname.find('BatchNorm1d') != -1: + init.normal_(m.weight.data, 1.0, 0.02) + init.constant_(m.bias.data, 0.0) + + +def weights_init_classifier(m): + classname = m.__class__.__name__ + if classname.find('Linear') != -1: + init.normal_(m.weight.data, std=0.001) + init.constant_(m.bias.data, 0.0) + + +# Defines the new fc layer and classification layer +# |--Linear--|--bn--|--relu--|--Linear--| +class ClassBlock(nn.Module): + def __init__(self, input_dim, class_num=1, activ='sigmoid', num_bottleneck=512): + super(ClassBlock, self).__init__() + + add_block = [] + add_block += [nn.Linear(input_dim, num_bottleneck)] + add_block += [nn.BatchNorm1d(num_bottleneck)] + add_block += [nn.LeakyReLU(0.1)] + add_block += [nn.Dropout(p=0.5)] + + add_block = nn.Sequential(*add_block) + add_block.apply(weights_init_kaiming) + + classifier = [] + classifier += [nn.Linear(num_bottleneck, class_num)] + if activ == 'sigmoid': + classifier += [nn.Sigmoid()] + elif activ == 'softmax': + classifier += [nn.Softmax()] + elif activ == 'none': + classifier += [] + else: + raise AssertionError("Unsupported activation: {}".format(activ)) + classifier = nn.Sequential(*classifier) + classifier.apply(weights_init_classifier) + + self.add_block = add_block + self.classifier = classifier + + def forward(self, x): + x = self.add_block(x) + x = self.classifier(x) + return x +