def imagesfeature():
coco_train = dset.CocoDetection(
root='E:/datasets/coco2017/train2017',
annFile='E:/datasets/coco2017/annotations/instances_train2017.json',
transform=transforms.ToTensor())
print('Number of samples: ', len(coco_train))
imgs = torch.zeros(1000, 3, 224, 224)
feats = np.zeros(shape=(len(coco_train), 4096))
ind = 0
alexnet = ImgNet()
alexnet.cuda()
for i in tqdm(range(len(coco_train))):
img, _ = coco_train[i]
img = np.array(img)
img = img.transpose((1, 2, 0))
img = reshapeImage(img)
imgs[ind, :] = img
if ind % 999 == 0 and ind is not 0:
feats[i - 999:i + 1, :] = alexnet(
imgs.cuda()).cpu().detach().numpy()
ind = 0
else:
ind = ind + 1
scio.savemat('cocomat/feature.mat', {'feature': feats})
print('finish!')
def test_draw_detections(writer):
print("Drawing some detections to tensorboard")
from utils.data_mappings import coco_id_to_name
from torchvision import datasets
from torchvision import transforms
def convert_bbox_labels(labels):
""" convert bounding boxes from left-upper-width-height format to
left-upper-right-down. """
for label in labels:
x, y, w, h = labels['bbox']
label['bbox'] = [x, y, x + w, y + h]
return labels
dataset = datasets.CocoDetection(
'../coco/images/val2017',
'../coco/annotations/instances_val2017.json',
transform=None,
target_transform=convert_bbox_labels
)
img, labels = dataset[0]
label_bboxs = [label['bbox'] for label in labels]
label_names = [coco_id_to_name[label['category_id']] for label in labels]
img = draw_detections(img, label_bboxs, label_names, scale=2.0)
writer.add_image("test_draw_detections", transforms.functional.to_tensor(img))
def createonehotlabelmat():
coco_train = dset.CocoDetection(
root='E:/datasets/coco2017/train2017',
annFile='E:/datasets/coco2017/annotations/instances_train2017.json',
transform=transforms.ToTensor())
print('Number of samples: ', len(coco_train))
onehot = np.zeros((len(coco_train), 91))
for i in tqdm(range(len(coco_train))):
_, target = coco_train[i]
labels = []
for t in target:
label = t['category_id']
# 去除重复的标签
if label not in labels:
onehot[i][label - 1] = 1
# scio.savemat('cocomat/labels.mat', {'labels': onehot})
matfile = 'cocomat/coco.mat'
temp = scio.loadmat(matfile)
scio.savemat(matfile, {
'T_bow1': temp['T_bow1'],
'T_bow2': temp['T_bow2'],
'L': onehot
})
print('create labels finish!')
def test_add_bbox():
from pycocotools.coco import COCO
detections = dset.CocoDetection(root=config.COCOData["train_img"],
annFile=config.COCOData["train_instances"],
transform=transforms.ToTensor())
print('Number of samples: ', len(detections))
print("Ploting image and bounding boxes...")
plt.show()
for img, targets in detections:
img = img.numpy().transpose((1, 2, 0))
img = np.ascontiguousarray(img)
print("showing image(" + str(img.shape) + ") with " +
str(len(targets)) + " bounding boxes")
gt_labels = []
for target in targets:
coordinates = target["bbox"]
gt_labels.append(target["category_id"])
print(coordinates)
x1, y1, x2, y2 = coor_transform(img.shape[0:2], coordinates)
img = cv2.rectangle(img, (x1, y1), (x2, y2),
color=(0, 255, 0),
thickness=2)
#cv2.rectangle(img, (0, 0), (20, 20), color=(0, 255, 0), thickness=3)
plt.imshow(img)
plt.pause(2)
def get_current_images_id(hparams, dataset_name):
"""
Get the list of image ids corresponding to break_training_loop_percentage in hparams.
It helps loading less captions, embeddings in create_list_of_captions_and_clean() , hence getting more samples per batches.
:param hparams: Hyperparameters
:param dataset_name: Dataset name to get images from
:return: List of image indices
"""
train_file = hparams[dataset_name]
image_dir = os.path.join(hparams['root'], train_file)
caption_file_path = get_cleaned_captions_path(hparams, train_file)
img_size = hparams["image_size"]
transform_pipeline = transforms.Compose([
CenteringPad(),
transforms.Resize((img_size, img_size)),
transforms.ToTensor()
])
coco_train_set = dset.CocoDetection(root=image_dir,
annFile=caption_file_path,
transform=transform_pipeline)
train_loader = torch.utils.data.DataLoader(
coco_train_set, batch_size=hparams["batch_size"])
break_training_loop_percentage = hparams["break_training_loop_percentage"]
break_training_loop_idx = max(
int(len(train_loader) * break_training_loop_percentage / 100) - 1, 0)
img_list = []
for idx, sample in enumerate(tqdm(train_loader)):
img_list.extend(sample[1][0]["image_id"].tolist())
if idx == break_training_loop_idx:
break
return img_list
def coco_dataset():
configs = yaml.safe_load(open('../config/coco.yml').read())
train_augmentation = transforms.Compose([ # transforms.RandomResizedCrop(800),#随机对图像进行裁剪，图像的标签会变化
# transforms.Resize(800), #只会根据短边修改到800，长边会变长
transforms.RandomCrop(800, padding=1, pad_if_needed=True), # 不够800填充黑色
transforms.RandomHorizontalFlip(0.5),
transforms.RandomVerticalFlip(0.5),
transforms.ToTensor(),
transforms.Normalize(configs['mean'], configs['std'])])
cap = det.CocoDetection(root='E:\\temp\\train_dataset_part1\\',
annFile='../gen/val_instance.json',
transform=train_augmentation) # 注意没有s
print('Number of sampler {}'.format(len(cap)))
for each in cap:
img, target = each ## load 4th sample
# print(img.size()) # 转化后图像
# print(target) # 标注,图像改变了大小，但是标注却没有？？？（不是学错了）
train_db, val_db = torch.utils.data.random_split(cap, [7, 3]) # 使用该方法会按照类别分类（放心使用）
train_loader = DataLoader(train_db, batch_size=5, shuffle=True)
val_loader = DataLoader(val_db, batch_size=1, shuffle=True)
for batch_index, data in enumerate(train_loader):
# data[0]存储的是原始img数据，len(data[0])=batch_size
# data[1]存储时原始标注信息
print(batch_index, len(data[0]), len(data[1]))
print(len(data[1]))
break
def create_dataset(dataset_type, dataset_path, transforms_para, batch_size,
is_debug=False, distributed=False, workers=4):
if dataset_type == DatasetsFactory.IMAGENET:
if is_debug:
train_dir = os.path.join(dataset_path, 'train_mini')
else:
train_dir = os.path.join(dataset_path, 'train')
val_dir = os.path.join(dataset_path, 'val')
train_dataset = datasets.ImageFolder(
train_dir,
transforms_para)
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=batch_size, shuffle=(train_sampler is None),
num_workers=workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(val_dir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(256),
transforms.ToTensor(),
])),
batch_size=32, shuffle=False,
num_workers=workers, pin_memory=True)
return train_loader, val_loader, train_sampler
if dataset_type == DatasetsFactory.COCO2014:
return datasets.CocoDetection()
logging.error("The dataset is not supported!")
return None
def get_coco_generator(img_path,
ann_path,
batch_size,
device,
shuffle=True,
isTrain=True):
transform = transforms.Compose(
[transforms.Resize((300, 300)),
transforms.ToTensor()])
def detection_collate_fn(sample_list):
img_batch = []
target_batch = []
for (x, y) in sample_list:
if not isTrain:
x_scale = 300 / x.size[0]
y_scale = 300 / x.size[1]
x = transform(x)
y = torch.stack([
y[:, 0], y[:, 1] * x_scale, y[:, 2] * y_scale,
y[:, 3] * x_scale, y[:, 4] * y_scale
],
dim=1)
else:
x, y = augment(x, y)
img_batch.append(x)
target_batch.append(y)
img_batch = torch.stack(img_batch).to(device)
return img_batch, target_batch
def coco_target_transform(y):
truths = []
for elem in y:
box = elem["bbox"]
truth = torch.zeros((5, ), device=device, dtype=torch.float)
truth[0] = elem["category_id"]
truth[1] = box[0] + 0.5 * box[2]
truth[2] = box[1] + 0.5 * box[3]
truth[3] = box[2]
truth[4] = box[3]
truths.append(truth)
if truths:
truth_array = torch.stack(truths, dim=0)
else:
truth_array = torch.zeros((0, 5), device=device, dtype=torch.float)
return truth_array
coco_data = dset.CocoDetection(root=img_path,
annFile=ann_path,
target_transform=coco_target_transform)
coco_generator = DataLoader(coco_data,
batch_size=batch_size,
collate_fn=detection_collate_fn,
shuffle=shuffle,
drop_last=True)
return coco_generator
def __init__(self, root, annotation_file, train=True):
self.train = train
self.coco = COCO(annotation_file)
self.dataset = datasets.CocoDetection(root,
annotation_file,
transform=None,
target_transform=None)
self.cur = 0
def __init__(self):
yolo_trans = transforms.Compose(
[transforms.Resize((416, 416)),
transforms.ToTensor()])
self.data = datasets.CocoDetection(
root="data/resized_val2017",
annFile="data/annotations/instances_val2017.json",
transform=yolo_trans)
self.data_iter = iter(self.data)
print("YoloGenInputActor init")
def __init__(self, COCO_path, is_training=True):
self.is_training = is_training
if is_training:
self.COCO_path = COCO_path + "train2017/"
self.coco_dataset = dset.CocoDetection(
root=self.COCO_path,
annFile=COCO_path + "annotations/instances_train2017.json")
else:
self.COCO_path = COCO_path + "val2017/"
self.coco_dataset = dset.CocoDetection(
root=self.COCO_path,
annFile=COCO_path + "annotations/instances_val2017.json")
self.unseen_classes = [
'cow', 'sheep', 'cat', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'chair', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sofa', 'train', 'tvmonitor'
]
def retinanet_coco_test(imgs, annotations, score_threshold, limit):
'''
Test pre-trained PyTorch RetinaNet against COCO.
Loads up a RetinaNet, runs data through it according to the provided annotation file
and evaluates the results using pycocotools COCOeval.
I used this to make sure that the PyTorch-provided RetinaNet works fine
and to familiarize myself with it!
'''
model = tmodels.detection.retinanet_resnet50_fpn(pretrained=True).cuda()
model.eval()
data = dsets.CocoDetection(root=imgs,
annFile=annotations,
transform=tforms.ToTensor())
evaluated = []
results = []
for i, datum in enumerate(data):
if i % 100 == 0:
print(i)
if limit >= 0 and i >= limit:
break
img, anns = datum
predictions = model(img.unsqueeze(0).cuda())[0]
keep = predictions['scores'] > score_threshold
boxes = [[
x1, y1, x2 - x1, y2 - y1
] for x1, y1, x2, y2 in utils.recall_tensor(predictions['boxes'][keep])
]
scores = utils.recall_tensor(predictions['scores'][keep])
labels = utils.recall_tensor(predictions['labels'][keep])
img_id = data.ids[i]
results += [{
'image_id': img_id,
'bbox': box,
'score': score,
'category_id': label
} for box, score, label in zip(boxes, scores, labels)]
evaluated.append(img_id)
evl = cocoeval.COCOeval(data.coco, data.coco.loadRes(results), 'bbox')
evl.params.imgIds = evaluated
evl.evaluate()
evl.accumulate()
evl.summarize()
ann_boxes = [ann['bbox'] for ann in anns]
ann_labels = [ann['category_id'] for ann in anns]
utils.show(img, boxes, ann_boxes, labels, ann_labels)
def get_coco_data_loader(data_path,
annotation_path,
batch_size,
train_split=None):
trfm = transforms.Compose(
[transforms.Resize((480, 640)),
transforms.ToTensor()])
data = datasets.CocoDetection(root=data_path,
annFile=annotation_path,
transform=trfm)
data = MultiviewDataset(data, [identity_view], [get_coco_captions])
return DataLoader(data, batch_size=batch_size)
def load_coco_val(config, args, hyper, distributed):
"""Validation set preprocessing and loader
"""
transform = transforms.Compose([transforms.ToTensor()])
valset = datasets.CocoDetection(config.val_path,
annFile=config.annFile,
transform=transform)
return torch.utils.data.DataLoader(valset,
batch_size=hyper.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
def load_dataset(root,
transform,
*args,
batch_size=32,
shuffle=True,
dataset_type='folder',
**kwargs):
"""
Parameters
-----------
dataset_type: str
should be voc , coco, cifar, minst or folder
if you're using voc dataset then you have to pass a param as year = 2007 or 2012
if you're using coco dataset then you have to pass a param as type = 'detection' or 'caption'
Return
----------
data: Dataloader
dataset: torchvision.dataset
"""
if dataset_type == 'folder':
dataset = datasets.ImageFolder(root, transform=transform)
elif dataset_type == 'voc':
year = kwargs.get('year', 2007)
image_set = kwargs.get('image_set', 'train')
dataset = datasets.VOCDetection(root,
year=year,
image_set=image_set,
transform=transform)
elif dataset_type == 'coco':
assert 'type' in kwargs and 'annfile' in kwargs
annfile = kwargs['annfile']
type = kwargs['type']
if type == 'detection':
dataset = datasets.CocoDetection(root,
annFile=annfile,
transform=transform)
elif type == 'caption':
dataset = datasets.CocoCaptions(root,
annFile=annfile,
transform=transform)
data = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
return data, dataset
def __init__(self,
conf,
path,
anno_path,
maps,
rotate=30,
shear=30,
PerspectiveTransform=0.1):
self.orig_dataset = datasets.CocoDetection(path, anno_path)
self.maps = maps
self.aug = Aug_Yolo(rotate, shear, PerspectiveTransform)
self.conf = conf
self.pair = namedtuple('pair', ['imgs', 'bboxes', 'labels'])
self.final_round = False
def __init__(self, is_train=True, transform=None):
if (is_train):
self.root = './datasets/coco/train2017'
self.annFile = './datasets/coco/annotations/person_keypoints_train2017.json'
else:
self.root = './datasets/coco/val2017'
self.annFile = './datasets/coco/annotations/person_keypoints_val2017.json'
if (transform == None):
self.transform = transforms.ToTensor()
else:
self.transform = transform
self.cap = dset.CocoDetection(root=self.root,
annFile=self.annFile,
transform=self.transform)
self.num = len(self.cap)
def imagesresize():
coco_train = dset.CocoDetection(
root='E:/datasets/coco2017/train2017',
annFile='E:/datasets/coco2017/annotations/instances_train2017.json',
transform=transforms.ToTensor())
print('Number of samples: ', len(coco_train))
imgs = np.zeros((len(coco_train), 3, 224, 224), dtype=np.uint8)
for i in tqdm(range(len(coco_train))):
img, _ = coco_train[i]
img = np.array(img)
img = img.transpose((1, 2, 0))
img = reshapeImage(img)
imgs[i, :] = img
# scio.savemat('cocomat/imageresize224.mat', {'': imgs})
f = h5py.File("cocomat/img-size224-numall-value127.hdf5", "w")
d1 = f.create_dataset("IAll", data=imgs)
print('finish!')
def createlabelmat():
coco_train = dset.CocoDetection(
root='E:/datasets/coco2017/train2017',
annFile='E:/datasets/coco2017/annotations/instances_train2017.json',
transform=transforms.ToTensor())
print('Number of samples: ', len(coco_train))
result = []
for i in tqdm(range(len(coco_train))):
_, target = coco_train[i]
labels = []
for t in target:
label = t['category_id']
# 去除重复的标签
if label not in labels:
labels.append(label)
result.append(labels)
scio.savemat('cocomat/labels.mat', {'labels': result})
print('create labels finish!')
def createlabeltxt():
coco_train = dset.CocoDetection(
root='E:/datasets/coco2017/train2017',
annFile='E:/datasets/coco2017/annotations/instances_train2017.json',
transform=transforms.ToTensor())
f = open('cocotxt/labels.txt', 'w')
for i in tqdm(range(len(coco_train))):
_, target = coco_train[i]
labels = []
for t in target:
label = t['category_id']
# 去除重复的标签
if label not in labels:
labels.append(label)
f.write(str(labels).replace('[', '').replace(']', '').replace(' ', ''))
f.write('\n')
labels.clear()
f.close()
print('create labels txt finish!')
def preprocess(dataset, annotations, file_prefix, size=256, num=-1):
coco_val = dset.CocoDetection(root=dataset,
annFile=annotations,
transform=transforms.ToTensor())
coco = COCO(annotation_file=annotations)
pk_images = []
pk_masks = []
num_to_process = num
if num_to_process == -1:
num_to_process = len(coco_val)
for i in range(num_to_process):
img, anns = coco_val[i]
# reshape img to be a tensor with dimensions [3, size, size] and add it to the image list
img = img.unsqueeze(0)
img = torch.nn.functional.interpolate(img, size=size)
img = img.squeeze(0)
pk_images.append(img)
# we create masks as a tensor with dimensions [91, size, size]. layer i represents the mask from the ith category
# WARNING: ALL MASKS OF THE SAME CATEGORY WILL BE COMBINED
# if the ith category is not present, then it will be all 0s
# for the categories that don't correspond to any object, their respective layers will also be 0
masks = torch.zeros((91, size, size))
for j in range(len(anns)):
mask = coco.annToMask(anns[j])
mask = torch.from_numpy(mask).unsqueeze(0)
mask = mask.unsqueeze(0)
mask = torch.nn.functional.interpolate(mask, size=size)
mask = mask.squeeze(0)
masks[int(anns[j]['category_id']), :, :] += mask.squeeze(0)
pk_masks.append(masks)
with open(file_prefix + "_images.pk", 'wb') as file:
pk.dump(pk_images, file)
with open(file_prefix + "_masks.pk", "wb") as file:
pk.dump(pk_masks, file)
return pk_images, pk_masks
def load_coco_data(config, args, hyper, distributed):
"""Training set preprocessing and loader
"""
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.CocoDetection(root=config.train_path,
annFile=config.annFile,
transform=transform)
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
dataset)
else:
train_sampler = None
return torch.utils.data.DataLoader(dataset,
batch_size=hyper.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
def main():
normalize = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = T.Compose(
[T.Resize(256),
T.CenterCrop(224),
T.ToTensor(), normalize])
testset = datasets.CocoDetection(
root="./cocoapi/images",
annFile="./cocoapi/annotations/instances_val2017.json",
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=1,
shuffle=False,
pin_memory=True)
resnet = models.resnet18(pretrained=True)
criterion = nn.CrossEntropyLoss()
validate(testloader, resnet, criterion)
def test_torch_COCOapi():
from pycocotools.coco import COCO
# The code is based on https://pytorch.org/docs/stable/torchvision/datasets.html#coco
detections = dset.CocoDetection(root=config.COCOData["train_img"],
annFile=config.COCOData["train_instances"],
transform=transforms.ToTensor())
print('Number of samples: ', len(detections))
img, targets = detections[3] # load 4th sample
print("Image Size: ", img.size())
# dict_keys(['segmentation', 'area', 'iscrowd', 'image_id', 'bbox', 'category_id', 'id'])
print("Annotation keys: ", targets[2].keys())
#plt.imshow(img.numpy().transpose((1, 2, 0)))
#plt.show()
print("Bounding boxes coordinates: ", end="")
print(targets[0]["bbox"])
print("type of Bounding boxes coordinates: ", end="")
print(type(targets[0]["bbox"][1]))
print("type of category_id:", end="")
print(type(targets[0]["category_id"]))
def load_dataset(root,
transform,
batch_size=32,
shuffle=True,
dataset_type='folder',
*args,
**kwargs):
"""
param
dataset_type: str
should be voc , coco, cifar, minst or folder
"""
if dataset_type == 'folder':
dataset = datasets.ImageFolder(root, transform=transform)
elif dataset_type == 'voc':
year = kwargs['year']
image_set = kwargs['image_set']
dataset = datasets.VOCDetection(root,
year=year,
image_set=image_set,
transform=transform)
elif dataset_type == 'coco':
annfile = kwargs['annfile']
type = kwargs['type']
if type == 'detect':
dataset = datasets.CocoDetection(root,
annFile=annfile,
transform=transform)
elif type == 'caption':
dataset = datasets.CocoCaptions(root,
annFile=annfile,
transform=transform)
data = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
return data, dataset.classes, dataset.class_to_idx
def __init__(self,
image_dir,
anno_path,
input_transform=None,
labels_path=None):
self.coco = dset.CocoDetection(root=image_dir, annFile=anno_path)
with open('./data/coco/category.json', 'r') as load_category:
self.category_map = json.load(load_category)
self.input_transform = input_transform
self.labels_path = labels_path
self.labels = []
if self.labels_path:
self.labels = np.load(self.labels_path).astype(np.float64)
self.labels = (self.labels > 0).astype(np.float64)
else:
l = len(self.coco)
for i in range(l):
item = self.coco[i]
print(i)
categories = self.getCategoryList(item[1])
label = self.getLabelVector(categories)
self.labels.append(label)
def visualize_coco_gaussians(imgs, annotations, method):
gauss_methods = {
'normal': {
'generate_method': datautils.generate_via_multivariate_normal(),
'join_method': datautils.join_via_max
},
'kant': {
'generate_method': datautils.generate_via_kant_method(),
'join_method': datautils.join_via_replacement
},
}
data = dsets.CocoDetection(root=imgs,
annFile=annotations,
transform=tforms.ToTensor())
img, anns = random.choice(data)
utils.show(img, groundtruth=[ann['bbox'] for ann in anns])
_, h, w = img.shape
coco_to_retina = lambda bbox: torch.tensor(
[bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]])
gauss = datautils.generate_gaussians(
w, h, [coco_to_retina(ann['bbox']) for ann in anns],
**gauss_methods[method])
utils.show(gauss)
def main():
args = get_args()
print('args:', args)
utils.init_distributed_mode(args)
# seed fixed
random_seed = args.seed
if args.seed is not None:
torch.manual_seed(random_seed)
torch.cuda.manual_seed(random_seed)
torch.cuda.manual_seed_all(random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(random_seed)
random.seed(random_seed)
warnings.warn(
'You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting from checkpoints.'
)
assert args.crop_size[0] <= args.train_size[0] and args.crop_size[1] <= args.train_size[1], \
'Mst be Crop size<= Image Size'
# create directory
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path + '/images')
if not os.path.isdir(args.save_path + '/results_color_val'):
os.makedirs(args.save_path + '/results_color_val')
os.makedirs(args.save_path + '/results_color_test')
# get sampler for loader considering distributed mode
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
dataset)
test_sampler = torch.utils.data.distributed.DistributedSampler(
dataset_test)
else:
train_sampler = torch.utils.data.RandomSampler(dataset)
test_sampler = torch.utils.data.SequentialSampler(dataset_test)
if args.aspect_ratio_group_factor >= 0:
group_ids = create_aspect_ratio_groups(
dataset, k=args.aspect_ratio_group_factor)
train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids,
args.batch_size)
else:
train_batch_sampler = torch.utils.data.BatchSampler(train_sampler,
args.batch_size,
drop_last=True)
# dataset selection
coco_train = datasets.CocoDetection(
root='/data/coco/images/train2017',
annFile='/data/coco/annotations/instances_train2017.json',
transform=get_transform(args))
coco_train_loader = DataLoader(coco_train,
batch_size=args.batch_size,
shuffle=True,
pin_memory=args.pin_memory,
num_workers=args.num_workers)
# collate_fn = collate_fn
# collate_fn = lambda x:x
coco_val = datasets.CocoDetection(
root='/data/coco/images/val2017',
annFile='/data/coco/annotations/instances_val2017.json',
transform=get_transform(args))
coco_val_loader = DataLoader(coco_train,
batch_size=args.batch_size,
shuffle=False,
pin_memory=args.pin_memory,
num_workers=args.num_workers)
# model, loss, optimizer
criterion = nn.CrossEntropyLoss()
breakpoint()
model = get_model(args)
model.cuda() # model.to(torch.device(args.device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
mode, device_ids=[args.gpu])
model_without_ddp = model.module
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr_init,
momentum=args.lr_momentum,
weight_decay=args.lr_weight_decay)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=args.epochs)
# Initialize metrics
best_miou = 0.0
metrics = {
'train_loss': [],
'train_acc': [],
'val_acc': [],
'val_loss': [],
'miou': []
}
start_epoch = 0
# resume training from checkpoint
if args.weights:
print('Resuming training from {}.'.format(args.weights))
checkpoint = torch.load(args.weights)
model.load_state_dict(checkpoint['model_state_dict'], strict=True)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
metrics = checkpoint['metrics']
best_miou = checkpoint['best_miou']
start_epoch = checkpoint['epoch'] + 1
# push model to GPU
# if torch.cuda.is_available():
# model = torch.nn.DataParallel(model).cuda()
# print('Model pushed to {} GPU(s), type {}.'.format(torch.cuda.device_count(), torch.cuda.get_device_name(0)))
# if predict mode
if args.predict:
checkpoint = torch.load(args.save_path + '/best_weights.pth.tar')
model.load_state_dict(checkpoint['model_state_dict'], strict=True)
print('Loaded model weights from {}'.format(args.save_path +
'/best_weights.pth.tar'))
# Create results directory
if not os.path.isdir(args.save_path + '/results_val'):
os.makedirs(args.save_path + '/results_val')
if not os.path.isdir(args.save_path + '/results_test'):
os.makedirs(args.save_path + '/results_test')
predict(coco_val_loader,
model,
folder=args.save_path,
mode='val',
args=args)
with open(args.save_path + '/log_epoch.csv', 'a') as epoch_log:
epoch_log.write(
'epoch, train loss, val loss, train acc, val acc, miou\n')
since = time.time()
for epoch in range(start_epoch, args.epochs):
# train
print('---training---')
train_loss, train_acc = train_epoch(coco_train_loader,
model=model,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=scheduler,
epoch=epoch,
args=args)
metrics['train_loss'].append(train_loss)
metrics['train_acc'].append(train_acc)
print('Epoch {} train loss: {:.4f}, acc: {:.4f}'.format(
epoch, train_loss, train_acc))
# validate
print('---validation---')
val_acc, val_loss, miou = validate_epoch(coco_val_loader,
model=model,
criterion=criterion,
epoch=epoch,
folder=args.save_path,
args=args)
metrics['val_acc'].append(val_acc)
metrics['val_loss'].append(val_loss)
metrics['miou'].append(miou)
# Write logs
with open(args.save_path + '/log_epoch.csv', 'a') as epoch_log:
epoch_log.write(
'{}, {:.5f}, {:.5f}, {:.5f}, {:.5f}, {:.5f}\n'.format(
epoch, train_loss, val_loss, train_acc, val_acc, miou))
# save ckpts
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'best_miou': best_miou,
'metrics': metrics,
}, args.save_path + '/checkpoint.pth.tar')
# Save best model to file
if miou > best_miou:
print('mIoU improved from {:.4f} to {:.4f}.'.format(
best_miou, miou))
best_miou = miou
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
}, args.save_path + '/best_weights.pth.tar')
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
plot_learning_curves(metrics, args)
# Load best model
checkpoint = torch.load(args.save_path + '/best_weights.pth.tar')
model.load_state_dict(checkpoint['model_state_dict'], strict=True)
print('Loaded best model weights (epoch {}) from {}/best_weights.pth.tar'.
format(checkpoint['epoch'], args.save_path))
# Create results directory
if not os.path.isdir(args.save_path + '/results_val'):
os.makedirs(args.save_path + '/results_val')
if not os.path.isdir(args.save_path + '/results_test'):
os.makedirs(args.save_path + '/results_test')
# Run prediction on validation set. For predicting on test set, simple replace 'val' by 'test'
predict(coco_val_loader,
model=model,
folder=args.save_path,
mode='val',
args=args)
import cv2
import random
font = cv2.FONT_HERSHEY_SIMPLEX
root = '/public/yzy/coco/2014/train2014/'
annFile = '/public/yzy/coco/2014/annotations/instances_train2014.json'
# 定义 coco collate_fn
def collate_fn_coco(batch):
return tuple(zip(*batch))
# 创建 coco dataset
coco_det = datasets.CocoDetection(root, annFile, transform=T.ToTensor())
# 创建 Coco sampler
sampler = torch.utils.data.RandomSampler(coco_det)
batch_sampler = torch.utils.data.BatchSampler(sampler, 8, drop_last=True)
# 创建 dataloader
data_loader = torch.utils.data.DataLoader(coco_det,
batch_sampler=batch_sampler,
num_workers=3,
collate_fn=collate_fn_coco)
# 可视化
for imgs, labels in data_loader:
for i in range(len(imgs)):
bboxes = []
ids = []
import torch
import torch.utils.data as data
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torch.autograd import Variable
batch_size = 1
transform = transforms.Compose([
transforms.Scale(600),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224,
0.225]),
])
det = dset.CocoDetection(root='data/train2014',
transform=transform,
annFile='data/annotations/instances_train2014.json')
loader = data.DataLoader(det,
batch_size=batch_size,
shuffle=True,
num_workers=6)
loader = iter(loader)
def load_model(model, name):
state_dict = torch.load('data/{}_state'.format(name))
model.load_state_dict(state_dict)
def next_batch():
# Retrieve next set of examples and initialize var for image data