ACAI/model-yolo-person-classify
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dk - edit

Browse Source
This commit is contained in:
CHO Daekeun
2026-07-09 18:16:06 +09:00
parent dda1759322
commit 4f1c320b3b
25 changed files with 1744 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4 -4
View File
@@ -66,7 +66,7 @@ clearml.conf
# etc
datafolder/
doc/
net/
test_sample/
# datafolder/
# doc/
# net/
# test_sample/
datafolder/folder.py
+140
View File
@@ -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
datafolder/reid_dataset/__init__.py
+17
View File
@@ -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
datafolder/reid_dataset/cuhk03_to_image.py
+47
View File
@@ -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)))
datafolder/reid_dataset/gdrive_downloader.py
+37
View File
@@ -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)
datafolder/reid_dataset/import_CUHK01.py
+22
View File
@@ -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
datafolder/reid_dataset/import_CUHK03.py
+46
View File
@@ -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
datafolder/reid_dataset/import_DukeMTMC.py
+32
View File
@@ -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
datafolder/reid_dataset/import_DukeMTMCAttribute.py
+126
View File
@@ -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
datafolder/reid_dataset/import_Market1501.py
+30
View File
@@ -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
datafolder/reid_dataset/import_Market1501Attribute.py
+174
View File
@@ -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
datafolder/reid_dataset/import_MarketDuke.py
+32
View File
@@ -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
datafolder/reid_dataset/import_MarketDuke_nodistractors.py
+33
View File
@@ -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
datafolder/reid_dataset/import_VIPeR.py
+29
View File
@@ -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
datafolder/reid_dataset/marketduke_to_hdf5.py
+32
View File
@@ -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()
datafolder/reid_dataset/pytorch_prepare.py
+96
View File
@@ -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)
datafolder/reid_dataset/reiddataset_downloader.py
+88
View File
@@ -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)
doc/DukeMTMC-reID-attribute
+16
View File
@@ -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.
doc/Market1501-attribute
+21
View File
@@ -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.
doc/attribute.json
+377
View File
@@ -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"
]
]
}
}
doc/label.json
+59
View File
@@ -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"
]
}
doc/save_json.py
+141
View File
@@ -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)
net/__init__.py
+9
View File
@@ -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)
net/models.py
+76
View File
@@ -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
net/utils.py
+60
View File
@@ -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