def main():
epoch = 9
test_batch_size = 64
scene_threshold = 0.
total_threshold = 0.2
num_of_vehicles = 64
cfg = get_default_config()
dataset = CityFlowNLInferenceDataset(cfg, build_transforms(cfg),
num_of_vehicles)
loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=8)
uuids, nls = query(cfg)
if os.path.exists('results'):
shutil.rmtree('results')
os.mkdir('results')
extract_cache_features(cfg, epoch, loader, dataset, test_batch_size, uuids,
nls)
cfg.num_gpu = torch.cuda.device_count()
cfg.resume_epoch = 0
mp.spawn(test,
args=(cfg, loader, dataset, epoch, uuids, nls, scene_threshold,
total_threshold),
nprocs=cfg.num_gpu,
join=True)
def test_model(cfg, model=None, to_file=False):
if to_file and os.path.exists('./results.txt'):
os.remove('./results.txt')
transforms = build_transforms(is_train=False)
data_loader = build_data_loader(cfg.dataset,
1,
cfg.workers,
transforms,
is_train=False)
if model is None:
model = build_model(data_loader.dataset.num_of_classes)
model = model.to(device)
load_dict(cfg, model, cfg.load_name)
model.eval()
running_loss, running_corrects, running_all, running_all_correct = 0., 0., 0., 0.
with torch.no_grad():
running_loss, running_corrects, running_all, running_all_correct = 0., 0., 0., 0.
for batch_idx, (data, targets,
target_lengths) in enumerate(data_loader):
data = data.to(device)
target_lengths = target_lengths.to(device)
predicted = model(data)
loss = loss_calculation(predicted,
torch.cat(targets, dim=0).to(device),
target_lengths)
running_loss += loss.item() * data.size(0)
running_corrects += get_accuracy(predicted, targets,
data_loader.dataset.ind_to_class,
to_file)
running_all_correct += torch.sum(target_lengths).item()
#print('sum : ', running_all_correct)
if to_file:
# print(data.cpu().numpy().squeeze().transpose(1,2,0))
# print(data.cpu()*255)
# torchvision.utils.save_image(data.cpu()*255, './test/{}.jpg'.format(batch_idx))
cv2.imwrite(
'./test/{}.jpg'.format(batch_idx),
data.cpu().numpy().squeeze().transpose(1, 2, 0) * 255)
running_all += len(data)
if batch_idx == 0:
since = time.time()
elif (batch_idx == len(data_loader) -
1): #batch_idx % cfg.interval == 0 or
print(
'Process: [{:5.0f}/{:5.0f} ({:.0f}%)]\tLoss: {:.4f}\tAcc:{:.4f}\tCost time:{:5.0f}s\tEstimated time:{:5.0f}s\r'
.format(running_all, len(data_loader.dataset),
100. * batch_idx / (len(data_loader) - 1),
running_loss / running_all,
running_corrects / running_all_correct,
time.time() - since, (time.time() - since) *
(len(data_loader) - 1) / batch_idx -
(time.time() - since))),
return running_corrects / running_all_correct
def train_model(cfg):
transforms = build_transforms(is_train=True)
data_loader = build_data_loader(cfg.dataset,
cfg.batch_size,
cfg.workers,
transforms,
is_train=True)
model = build_model(data_loader.dataset.num_of_classes)
model = model.to(device)
model.train()
optimizer = get_optimizer(cfg, model)
if cfg.start_epoch > 0:
load_dict(cfg, model, cfg.start_epoch - 1, optimizer)
best_accu = 0.
for epoch in range(cfg.start_epoch, cfg.epochs):
model.train()
running_loss, running_corrects, running_all, running_all_correct = 0., 0., 0., 0.
for batch_idx, (data, targets,
target_lengths) in enumerate(data_loader):
data = data.to(device)
target_lengths = target_lengths.to(device)
predicted = model(data)
loss = loss_calculation(predicted,
torch.cat(targets, dim=0).to(device),
target_lengths)
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
running_loss += loss.item() * data.size(0)
running_corrects += get_accuracy(
predicted, targets, data_loader.dataset.ind_to_class)
running_all_correct += torch.sum(target_lengths).item()
#print('sum : ', running_all_correct)
running_all += len(data)
if batch_idx == 0:
since = time.time()
elif batch_idx % cfg.interval == 0 or (batch_idx
== len(data_loader) - 1):
print(
'Process: [epoch {} {:5.0f}/{:5.0f} ({:.0f}%)]\tLoss: {:.4f}\tAcc:{:.4f}\tCost time:{:5.0f}s\tEstimated time:{:5.0f}s\r'
.format(epoch, running_all, len(data_loader.dataset),
100. * batch_idx / (len(data_loader) - 1),
running_loss / running_all,
running_corrects / running_all_correct,
time.time() - since, (time.time() - since) *
(len(data_loader) - 1) / batch_idx -
(time.time() - since))),
val_accu = test_model(cfg, model)
if val_accu > best_accu:
best_accu = val_accu
save_state(cfg, model, optimizer, 'best', best_accu)
save_state(cfg, model, optimizer, cfg.epochs, val_accu)
import cv2
import numpy as np
import os
import shutil
from tqdm import tqdm
from glob import glob
import json
epoch = 19
test_batch_size = 64
scene_threshold = 0.8
total_threshold = 0.5
num_of_vehicles = 64
cfg = get_default_config()
dataset = CityFlowNLInferenceDataset(cfg, build_transforms(cfg),
num_of_vehicles)
model = MyModel(cfg,
len(dataset.nl),
dataset.nl.word_to_idx['<PAD>'],
nn.BatchNorm2d,
num_colors=len(CityFlowNLDataset.colors),
num_types=len(CityFlowNLDataset.vehicle_type) - 2).cuda()
loader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=4)
uuids, nls = query(cfg)
saved_dict = torch.load(f'save/{epoch}.pth')
n = {}
for k, v in saved_dict.items():
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
assert (cfg.TEST.IMS_PER_BATCH == 1)
img_names = []
if not args.root_dir.endswith('/'):
args.root_dir += '/'
for root, parent, img_files in os.walk(args.root_dir):
for img_file in img_files:
if img_file.lower().endswith(('.jpg', '.png', '.bmp', '.jpeg')):
img_names.append(
os.path.join(root, img_file)[len(args.root_dir):])
transformer = build_transforms(cfg)
coco_demo = COCODemo(cfg)
start_time = time.time()
end_time = time.time()
for img_idx, img_name in enumerate(img_names[:100000]):
print('{}/{}: img_name={}'.format(img_idx, len(img_names), img_name))
img_ori = cv2.imread(os.path.join(args.root_dir, img_name))
img = transformer(img_ori[:, :, ::-1])
top_predictions = coco_demo.get_top_detections(img,
img_ori.shape[0],
img_ori.shape[1],
img_name,
topn=None)
print("using time: {:.4f}s".format(time.time() - end_time))
end_time = time.time()
print("Average time {:.4f}s ({:.4f}s/{})".format(
#print(len(target))
target_m = []
if len(target.get_field('masks')) > 1:
print(anno[0]['image_id'])
for m in target.get_field('masks'):
target_m.append(m.get_mask_tensor().float().unsqueeze(0))
masks_target = torch.cat(target_m, dim=0)
masks_target = (torch.sum(masks_target, dim=0) > 0).float().unsqueeze(0)
return img, masks_target, [ori_size, self.ids[idx]]
if __name__=='__main__':
import cv2
import sys
from transforms import build_transforms
df = '/core1/data/home/niuwenhao/workspace/data/detection/door_all_new.json'
cd = COCODataset(df, '', True, transforms=build_transforms(),clamp=False)
for im, targets, idx in cd:
sys.stdout.write('{} / {}\r'.format(idx, len(cd)))
sys.stdout.flush()
continue
print(im.shape, targets.get_field('masks').get_mask_tensor().shape)
for m in targets.get_field('masks'):
show = 255 * im * m.get_mask_tensor()
show = show.numpy().astype(np.uint8)
print(show)
print(show.shape)
print(show.dtype)
cv2.imshow('test', show)
cv2.waitKey()
def __init__(self,
root,
training=True,
crop_size=(512, 1024),
mirror=True,
min_scale=0.5,
max_scale=2.,
scale_step_size=0.25,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
semantic_only=False,
ignore_stuff_in_offset=False,
small_instance_area=0,
small_instance_weight=1):
"""Cityscapes panoptic segmentation dataset.
Arguments:
root: Str, root directory.
training: Bool, for training or testing.
crop_size: Tuple, crop size. (height, width)
mirror: Bool, whether to apply random horizontal flip.
min_scale: Float, min scale in scale augmentation.
max_scale: Float, max scale in scale augmentation.
scale_step_size: Float, step size to select random scale.
mean: Tuple, image mean.
std: Tuple, image std.
semantic_only: Bool, only use semantic segmentation label.
ignore_stuff_in_offset: Boolean, whether to ignore stuff region when training the offset branch.
small_instance_area: Integer, indicates largest area for small instances.
small_instance_weight: Integer, indicates semantic loss weights for small instances.
"""
self.root = root
self.training = training # If training=False, target is None.
self.crop_h, self.crop_w = crop_size
self.mirror = mirror
self.min_scale = min_scale
self.max_scale = max_scale
self.scale_step_size = scale_step_size
self.mean = mean
self.std = std
self.pad_value = tuple([int(v * 255) for v in self.mean])
self.ignore_label = 0
self.label_pad_value = (0, 0, 0)
# Define instance variable
self.rgb2id = rgb2id
self.id2rgb = id2rgb
self.image_files = [] # list of file names
self.image_elements = [] # image element in xml file
for d in sorted(os.listdir(self.root)):
# collect image files
files = sorted(glob(os.path.join(self.root, d, '*.jpg')))
self.image_files += files
if self.training:
# read xml files and collect image elements
xml_path = glob(os.path.join(self.root, d, '*.xml'))[0]
root_element = elemTree.parse(xml_path)
image_elements = root_element.findall('./image')
self.image_elements += image_elements
# Define transforms
self.transform = build_transforms(self, training)
if self.training:
if semantic_only:
self.target_transform = SemanticTargetGenerator(
self.ignore_label, self.rgb2id)
else:
self.target_transform = PanopticTargetGenerator(
self.ignore_label,
self.rgb2id,
_ROAD_CONDITION_THINGS_LIST,
sigma=8,
ignore_stuff_in_offset=ignore_stuff_in_offset,
small_instance_area=small_instance_area,
small_instance_weight=small_instance_weight)
else:
self.target_transform = None
def train_model_on_dataset(rank, cfg):
dist_rank = rank
# print(dist_rank)
dist.init_process_group(backend="nccl", rank=dist_rank,
world_size=cfg.num_gpu,
init_method="env://")
torch.cuda.set_device(rank)
cudnn.benchmark = True
dataset = CityFlowNLDataset(cfg, build_transforms(cfg))
model = MyModel(cfg, len(dataset.nl), dataset.nl.word_to_idx['<PAD>'], norm_layer=nn.SyncBatchNorm, num_colors=len(CityFlowNLDataset.colors), num_types=len(CityFlowNLDataset.vehicle_type) - 2).cuda()
model = DistributedDataParallel(model, device_ids=[rank],
output_device=rank,
broadcast_buffers=cfg.num_gpu > 1, find_unused_parameters=False)
optimizer = torch.optim.Adam(
params=model.parameters(),
lr=cfg.TRAIN.LR.BASE_LR, weight_decay=0.00003)
lr_scheduler = WarmupMultiStepLR(optimizer,
milestones=cfg.TRAIN.STEPS,
gamma=cfg.TRAIN.LR.WEIGHT_DECAY,
warmup_factor=cfg.TRAIN.WARMUP_FACTOR,
warmup_iters=cfg.TRAIN.WARMUP_EPOCH)
color_loss = LabelSmoothingLoss(len(dataset.colors), 0.1)
vehicle_loss = LabelSmoothingLoss(len(dataset.vehicle_type) - 2, 0.1)
if cfg.resume_epoch > 0:
model.load_state_dict(torch.load(f'save/{cfg.resume_epoch}.pth'))
optimizer.load_state_dict(torch.load(f'save/{cfg.resume_epoch}_optim.pth'))
lr_scheduler.last_epoch = cfg.resume_epoch
lr_scheduler.step()
if rank == 0:
print(f'resume from {cfg.resume_epoch} pth file, starting {cfg.resume_epoch+1} epoch')
cfg.resume_epoch += 1
# loader = DataLoader(dataset, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=True, num_workers=cfg.TRAIN.NUM_WORKERS)
train_sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=cfg.TRAIN.BATCH_SIZE //cfg.num_gpu,
num_workers=cfg.TRAIN.NUM_WORKERS // cfg.num_gpu,# shuffle=True,
sampler=train_sampler, pin_memory=True)
for epoch in range(cfg.resume_epoch, cfg.TRAIN.EPOCH):
losses = 0.
losses_color = 0.
losses_types = 0.
losses_nl_color = 0.
losses_nl_types = 0.
precs = 0.
train_sampler.set_epoch(epoch)
for idx, (nl, frame, label, act_map, color_label, type_label, nl_color_label, nl_type_label) in enumerate(loader):
# print(nl.shape)
# print(global_img.shape)
# print(local_img.shape)
nl = nl.cuda(non_blocking=True)
label = label.cuda(non_blocking=True)
act_map = act_map.cuda(non_blocking=True)
# global_img, local_img = global_img.cuda(), local_img.cuda()
nl = nl.transpose(1, 0)
frame = frame.cuda(non_blocking=True)
color_label = color_label.cuda(non_blocking=True)
type_label = type_label.cuda(non_blocking=True)
nl_color_label = nl_color_label.cuda(non_blocking=True)
nl_type_label = nl_type_label.cuda(non_blocking=True)
output, color, types, nl_color, nl_types = model(nl, frame, act_map)
# loss = sampling_loss(output, label, ratio=5)
# loss = F.binary_cross_entropy_with_logits(output, label)
total_num_pos = reduce_sum(label.new_tensor([label.sum()])).item()
num_pos_avg_per_gpu = max(total_num_pos / float(cfg.num_gpu), 1.0)
loss = sigmoid_focal_loss(output, label, reduction='sum') / num_pos_avg_per_gpu
loss_color = color_loss(color, color_label) * cfg.TRAIN.ALPHA_COLOR
loss_type = vehicle_loss(types, type_label) * cfg.TRAIN.ALPHA_TYPE
loss_nl_color = color_loss(nl_color, nl_color_label) * cfg.TRAIN.ALPHA_NL_COLOR
loss_nl_type = vehicle_loss(nl_types, nl_type_label) * cfg.TRAIN.ALPHA_NL_TYPE
loss_total = loss + loss_color + loss_type + loss_nl_color + loss_nl_type
optimizer.zero_grad()
loss_total.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
optimizer.step()
losses += loss.item()
losses_color += loss_color.item()
losses_types += loss_type.item()
losses_nl_color += loss_nl_color.item()
losses_nl_types += loss_nl_type.item()
# precs += recall.item()
if rank == 0 and idx % cfg.TRAIN.PRINT_FREQ == 0:
pred = (output.sigmoid() > 0.5)
# print((pred == label).sum())
pred = (pred == label)
recall = (pred * label).sum() / label.sum()
ca = (color.argmax(dim=1) == color_label)
ca = ca.sum().item() / ca.numel()
ta = (types.argmax(dim=1) == type_label)
ta = ta.sum().item() / ta.numel()
# accu = pred.sum().item() / pred.numel()
lr = optimizer.param_groups[0]['lr']
print(f'epoch: {epoch},',
f'lr: {lr}, step: {idx}/{len(loader)},',
f'loss: {losses / (idx + 1):.4f},',
f'loss color: {losses_color / (idx + 1):.4f},',
f'loss type: {losses_types / (idx + 1):.4f},',
f'loss nl color: {losses_nl_color / (idx + 1):.4f},',
f'loss nl type: {losses_nl_types / (idx + 1):.4f},',
f'recall: {recall.item():.4f}, c_accu: {ca:.4f}, t_accu: {ta:.4f}')
lr_scheduler.step()
if rank == 0:
if not os.path.exists('save'):
os.mkdir('save')
torch.save(model.state_dict(), f'save/{epoch}.pth')
torch.save(optimizer.state_dict(), f'save/{epoch}_optim.pth')