def build_model(cfg, args):
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_ssd_model(cfg)
return model
def run_demo(cfg, weights_file, iou_threshold, score_threshold, images_dir, output_dir, dataset_type):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
else:
raise NotImplementedError('Not implemented now.')
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg, is_test=True)
model.load(weights_file)
print('Loaded weights from {}.'.format(weights_file))
model = model.to(device)
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
output = predictor.predict(image)
boxes, labels, scores = [o.to(cpu_device).numpy() for o in output]
drawn_image = draw_bounding_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
image_name = os.path.basename(image_path)
Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))
def train(cfg, args):
#logging.basicConfig(filename='./output/LOG/'+__name__+'.log',format='[%(asctime)s-%(filename)s-%(levelname)s:%(message)s]', level = logging.INFO,filemode='a',datefmt='%Y-%m-%d%I:%M:%S %p')
logging.basicConfig(
filename='./output/LOG/' + __name__ + '.log',
format='[%(asctime)s-%(filename)s-%(levelname)s:%(message)s]',
level=logging.INFO)
logger = logging.getLogger('SSD.trainer')
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
ssd_model = build_ssd_model(cfg)
ssd_model.init_from_base_net(args.vgg)
ssd_model = torch.nn.DataParallel(ssd_model,
device_ids=range(
torch.cuda.device_count()))
device = torch.device(cfg.MODEL.DEVICE)
print(ssd_model)
logger.info(ssd_model)
model = torchvision.models.AlexNet(num_classes=10)
logger.info(model)
writer = tensorboardX.SummaryWriter(log_dir="./output/model_graph/",
comment="myresnet")
#dummy_input = torch.autograd.Variable(torch.rand(1, 3, 227, 227))
dummy_input = torch.autograd.Variable(torch.rand(1, 3, 300, 300))
#writer.add_graph(model=ssd_model, input_to_model=(dummy_input, ))
model_onnx_path = 'torch_model.onnx'
output = torch_onnx.export(ssd_model,
dummy_input,
model_onnx_path,
verbose=False)
#ssd_model.to(device)
print('----------------')
def evaluation(cfg, weights_file, output_dir, distributed):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
model.load(weights_file)
logger = logging.getLogger("SSD.inference")
logger.info('Loaded weights from {}.'.format(weights_file))
model.to(device)
do_evaluation(cfg, model, output_dir, distributed)
def convert2scriptmodule(cfg, args):
ssd_model = build_ssd_model(cfg)
print(ssd_model)
input = torch.Tensor(1, 3, cfg.INPUT.IMAGE_SIZE, cfg.INPUT.IMAGE_SIZE)
model_path = args.model_path # 'ssd300_vgg_final.pth'
ssd_model.load_state_dict(torch.load(model_path))
save(ssd_model, input, args.model_out) # 'script.pt'
def build_model(cfg, args):
cfg.merge_from_file("configs/ssd512_voc0712.yaml")
#cfg.merge_from_list(args.opts)
cfg.freeze()
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_ssd_model(cfg)
return model
def setup_self_ade(cfg, args):
logger = logging.getLogger("self_ade.setup")
logger.info("Starting self_ade setup")
# build model from config
model = build_ssd_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
train_transform = TrainAugmentation(cfg.INPUT.IMAGE_SIZE,
cfg.INPUT.PIXEL_MEAN,
cfg.INPUT.PIXEL_STD)
target_transform = MatchPrior(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE,
cfg.MODEL.THRESHOLD)
test_dataset = build_dataset(dataset_list=cfg.DATASETS.TEST,
is_test=True)[0]
self_ade_dataset = build_dataset(dataset_list=cfg.DATASETS.TEST,
transform=train_transform,
target_transform=target_transform)
ss_dataset = SelfSupervisedDataset(self_ade_dataset, cfg)
test_sampler = SequentialSampler(test_dataset)
os_sampler = OneSampleBatchSampler(test_sampler, cfg.SOLVER.BATCH_SIZE,
args.self_ade_iterations)
self_ade_dataloader = DataLoader(ss_dataset,
batch_sampler=os_sampler,
num_workers=args.num_workers)
effective_lr = args.learning_rate * args.self_ade_weight
optimizer = torch.optim.SGD(model.parameters(),
lr=effective_lr,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
# Initialize mixed-precision training
use_mixed_precision = cfg.USE_AMP
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
execute_self_ade(cfg, args, test_dataset, self_ade_dataloader, model,
optimizer)
def run_demo(cfg, weights_file, iou_threshold, score_threshold, video_dir, output_dir, dataset_type):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
elif dataset_type == 'cla':
class_names = CLADataset.class_names
else:
raise NotImplementedError('Not implemented now.')
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
model.load(weights_file)
print('Loaded weights from {}.'.format(weights_file))
model = model.to(device)
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
stream = cv2.VideoCapture(video_dir)
# image_paths = glob.glob(os.path.join(video_dir, '*.jpg'))
# 获得输出视频大小，与原视频大小相同
shape=(int(stream.get(cv2.CAP_PROP_FRAME_WIDTH)),int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT)))
# 获取输出视频的帧率
_fps = stream.get(cv2.CAP_PROP_FPS)
# 指定视频编码
fourcc = cv2.VideoWriter_fourcc(*'XVID')
# 如果输出目录不存在 创建输出目录
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 创建视频输出对象
output_name = os.path.basename(video_dir)
output_name = output_name.split('.')[0] + ".avi"
writer = cv2.VideoWriter(os.path.join(output_dir, output_name), fourcc, _fps, shape)
while True:
ret, image = stream.read()
if ret is False :
break
output = predictor.predict(image)
boxes, labels, scores = [o.to(cpu_device).numpy() for o in output]
drawn_image = draw_bounding_boxes(image, boxes, labels, scores, class_names).astype(np.uint8)
writer.write(drawn_image)
writer.release()
stream.release()
def load_model(config_file='SSD/configs/ssd300_voc0712.yaml',
iou_threshold=0.5,
score_threshold=0.5):
cfg = get_configuration(config_file)
class_names = VOCDataset.class_names
global device
device = torch.device('cpu')
model = build_ssd_model(cfg)
logger.info('Loading model from S3')
obj = s3.get_object(Bucket=MODEL_BUCKET, Key=MODEL_KEY)
bytestream = io.BytesIO(obj['Body'].read())
model.load(bytestream)
logger.info(f'Loaded weights from {MODEL_KEY}')
model = model.to(device)
return Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
def evaluation(cfg, args, weights_file, output_dir, distributed):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
model.load(open(weights_file, 'rb'))
logger = logging.getLogger("SSD.inference")
logger.info('Loaded weights from {}.'.format(weights_file))
model.to(device)
if args.eval_mode == "test":
do_evaluation(cfg, model, output_dir, distributed)
else:
dataset_metrics = do_evaluation(cfg, model, cfg.OUTPUT_DIR, distributed, datasets_dict=_create_val_datasets(args, cfg, logger))
count = len(dataset_metrics)
map_sum = 0
for k,v in dataset_metrics.items():
#logger.info("mAP on {}: {:.3f}".format(k, v.info["mAP"]))
map_sum += v.info["mAP"]
avg_map = map_sum/count
print("'Model': '{}', 'Avg_mAP': {}".format(weights_file, avg_map))
def run_demo(cfg, checkpoint_file, iou_threshold, score_threshold, images_dir,
output_dir):
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['state_dict'])
print('Loaded weights from {}.'.format(checkpoint_file))
model = model.to(device)
model.eval()
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
add_count = 0
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
# image_mirror = image[:, ::-1]
output = predictor.predict(image)
boxes, scores, seg_map = [o.to(cpu_device).numpy() for o in output]
seg_map = cv2.resize(seg_map, (512, 512)) * 255
seg_map = seg_map.astype(np.uint8)
# seg_map = cv2.applyColorMap(seg_map, cv2.COLORMAP_JET)
seg_map = cv2.resize(seg_map, (1280, 720),
interpolation=cv2.INTER_CUBIC)
drawn_image = draw_bounding_boxes(image, boxes).astype(np.uint8)
image_name = os.path.basename(image_path)
txt_path = os.path.join(output_dir, 'txtes')
if not os.path.exists(txt_path):
os.makedirs(txt_path)
txt_path = os.path.join(txt_path,
'res_' + image_name.replace('jpg', 'txt'))
#multi-output merge
merge_output = False
if merge_output:
ret, binary = cv2.threshold(seg_map, 75, 255, cv2.THRESH_BINARY)
# cv2.imshow('binary:',binary)
# cv2.waitKey()
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
w, h = np.shape(binary)
for contour in contours:
# 获取最小包围矩形
rect = cv2.minAreaRect(contour)
# 中心坐标
x, y = rect[0]
# cv2.circle(img, (int(x), int(y)), 3, (0, 255, 0), 5)
# 长宽,总有 width>=height
width, height = rect[1]
if width < 10 or height < 10:
continue
# 角度:[-90,0)
angle = rect[2]
box = cv2.boxPoints(rect)
box = np.int0(box)
box[:, 0] = np.clip(box[:, 0], 0, h)
box[:, 1] = np.clip(box[:, 1], 0, w)
poly1 = Polygon(box).convex_hull
intersect = False
for item in boxes:
print('item:', item)
poly2 = Polygon(item.reshape(4, 2)).convex_hull
if poly1.intersects(poly2): # 如果两四边形相交
intersect = True
break
if not intersect:
print('boxes.shape:', np.shape(boxes))
box = box.reshape((1, 8))
print('box.shape:', np.shape(box))
num, _ = np.shape(boxes)
if num == 0:
print('num == 0')
boxes = box
else:
boxes = np.concatenate((boxes, box))
print('boxes.shape:', np.shape(boxes))
print('add one box')
add_count += 1
# cv2.line(image, (box[0][0], box[0][1]), (box[0][2], box[0][3]), (0, 0, 255), thickness=4)
# cv2.line(image,(box[0][2], box[0][3]), (box[0][4], box[0][5]), (0, 0, 255), thickness=4)
# cv2.line(image,(box[0][4], box[0][5]), (box[0][6], box[0][7]), (0, 0, 255), thickness=4)
# cv2.line(image, (box[0][6], box[0][7]), (box[0][0], box[0][1]), (0, 0, 255), thickness=4)
# cv2.imshow('img',image)
# cv2.waitKey()
# print('txt_path:',txt_path)
with open(txt_path, 'w+') as f:
for box in boxes:
box_temp = np.reshape(box, (4, 2))
box = order_points_quadrangle(box_temp)
box = np.reshape(box, (1, 8)).squeeze(0)
is_valid = validate_clockwise_points(box)
if not is_valid:
continue
# print('box:',box)
line = ''
for item in box:
if item < 0:
item = 0
line += str(int(item)) + ','
line = line[:-1] + '\n'
f.write(line)
path = os.path.join(output_dir, image_name)
print('path:', path)
Image.fromarray(drawn_image).save(path)
path = os.path.join(
output_dir,
image_name.split('.')[0] + '_segmap.' + image_name.split('.')[1])
# print(path)
# 存储score_map
cv2.imwrite(path, seg_map)
print('add count:', add_count)
def train(cfg, args):
logger = logging.getLogger('SSD.trainer')
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_ssd_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if args.resume:
logger.info("Resume from the model {}".format(args.resume))
model.load(args.resume)
else:
logger.info("Init from base net {}".format(args.vgg))
model.init_from_base_net(args.vgg)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
# -----------------------------------------------------------------------------
# Optimizer
# -----------------------------------------------------------------------------
lr = cfg.SOLVER.LR * args.num_gpus # scale by num gpus
optimizer = torch.optim.SGD(model.parameters(),
lr=lr,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
# -----------------------------------------------------------------------------
# Criterion
# -----------------------------------------------------------------------------
criterion = MultiBoxLoss(iou_threshold=cfg.MODEL.THRESHOLD,
neg_pos_ratio=cfg.MODEL.NEG_POS_RATIO)
# -----------------------------------------------------------------------------
# Scheduler
# -----------------------------------------------------------------------------
milestones = [step // args.num_gpus for step in cfg.SOLVER.LR_STEPS]
scheduler = WarmupMultiStepLR(optimizer=optimizer,
milestones=milestones,
gamma=cfg.SOLVER.GAMMA,
warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
warmup_iters=cfg.SOLVER.WARMUP_ITERS)
# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
train_transform = TrainAugmentation(cfg.INPUT.IMAGE_SIZE,
cfg.INPUT.PIXEL_MEAN)
target_transform = MatchPrior(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE,
cfg.MODEL.THRESHOLD)
train_dataset = build_dataset(dataset_list=cfg.DATASETS.TRAIN,
transform=train_transform,
target_transform=target_transform)
logger.info("Train dataset size: {}".format(len(train_dataset)))
if args.distributed:
sampler = torch.utils.data.DistributedSampler(train_dataset)
else:
sampler = torch.utils.data.RandomSampler(train_dataset)
batch_sampler = torch.utils.data.sampler.BatchSampler(
sampler=sampler, batch_size=cfg.SOLVER.BATCH_SIZE, drop_last=False)
batch_sampler = samplers.IterationBasedBatchSampler(
batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER // args.num_gpus)
train_loader = DataLoader(train_dataset,
num_workers=4,
batch_sampler=batch_sampler)
return do_train(cfg, model, train_loader, optimizer, scheduler, criterion,
device, args)
net = add_flops_counting_methods(net)
net = net.cuda().eval()
net.start_flops_count()
_ = net(input)
return net.compute_average_flops_cost()/1e9/2
# example
if __name__ == '__main__':
from ssd.modeling.vgg_ssd import build_ssd_model
from ssd.config import cfg
'''
'''
cfg.merge_from_file("configs/ssd512_voc0712.yaml")
cfg.freeze()
model = build_ssd_model(cfg)
input_size = (1024, 1024)
#ssd_net = model.eval()
ssd_net = model.cuda()
total_flops = get_flops(ssd_net, input_size)
# For default vgg16 model, this shoud output 31.386288 G FLOPS
print("The Model's Total FLOPS is : {:.6f} G FLOPS".format(total_flops))
def train(cfg, args):
logger = logging.getLogger('SSD.trainer')
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_ssd_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
if args.resume:
logger.info("Resume from the model {}".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
iteration = checkpoint['iteration']
print('iteration:', iteration)
elif args.vgg:
iteration = 0
logger.info("Init from backbone net {}".format(args.vgg))
model.init_from_base_net(args.vgg)
else:
iteration = 0
logger.info("all init from kaiming init")
# -----------------------------------------------------------------------------
# Optimizer
# -----------------------------------------------------------------------------
lr = cfg.SOLVER.LR * args.num_gpus # scale by num gpus
#optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=cfg.SOLVER.MOMENTUM, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
print('cfg.SOLVER.WEIGHT_DECAY:', cfg.SOLVER.WEIGHT_DECAY)
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
# -----------------------------------------------------------------------------
# Scheduler
# -----------------------------------------------------------------------------
milestones = [step // args.num_gpus for step in cfg.SOLVER.LR_STEPS]
scheduler = WarmupMultiStepLR(optimizer=optimizer,
milestones=milestones,
gamma=cfg.SOLVER.GAMMA,
warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
warmup_iters=cfg.SOLVER.WARMUP_ITERS)
# ------------------------1-----------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
#对原始图像进行数据增强
train_transform = TrainAugmentation(cfg.INPUT.IMAGE_SIZE,
cfg.INPUT.PIXEL_MEAN)
target_transform = MatchPrior(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE,
cfg.MODEL.IOU_THRESHOLD, cfg.MODEL.PRIORS.DISTANCE_THRESHOLD)
train_dataset = build_dataset(dataset_list=cfg.DATASETS.TRAIN,
transform=train_transform,
target_transform=target_transform,
args=args)
logger.info("Train dataset size: {}".format(len(train_dataset)))
sampler = torch.utils.data.RandomSampler(train_dataset)
# sampler = torch.utils.data.SequentialSampler(train_dataset)
batch_sampler = torch.utils.data.sampler.BatchSampler(
sampler=sampler, batch_size=cfg.SOLVER.BATCH_SIZE, drop_last=False)
batch_sampler = samplers.IterationBasedBatchSampler(
batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER // args.num_gpus)
train_loader = DataLoader(train_dataset,
num_workers=4,
batch_sampler=batch_sampler,
pin_memory=True)
return do_train(cfg, model, train_loader, optimizer, scheduler, device,
args, iteration)
def train(cfg, args):
logger = logging.getLogger('SSD.trainer')
# -----------------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------------
model = build_ssd_model(cfg)
device = torch.device(cfg.MODEL.DEVICE)
model.to(device)
# -----------------------------------------------------------------------------
# Optimizer
# -----------------------------------------------------------------------------
lr = cfg.SOLVER.LR * args.num_gpus # scale by num gpus
optimizer = torch.optim.SGD(model.parameters(),
lr=lr,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY)
# -----------------------------------------------------------------------------
# Scheduler
# -----------------------------------------------------------------------------
milestones = [step // args.num_gpus for step in cfg.SOLVER.LR_STEPS]
scheduler = WarmupMultiStepLR(optimizer=optimizer,
milestones=milestones,
gamma=cfg.SOLVER.GAMMA,
warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
warmup_iters=cfg.SOLVER.WARMUP_ITERS)
# -----------------------------------------------------------------------------
# Load weights or restore checkpoint
# -----------------------------------------------------------------------------
if args.resume:
logger.info("Resume from the model {}".format(args.resume))
restore_training_checkpoint(logger,
model,
args.resume,
optimizer=optimizer,
scheduler=scheduler)
else:
logger.info("Init from base net {}".format(args.vgg))
model.init_from_base_net(args.vgg)
# Initialize mixed-precision training
use_mixed_precision = cfg.USE_AMP
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model, optimizer = amp.initialize(model,
optimizer,
opt_level=amp_opt_level)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
train_transform = TrainAugmentation(cfg.INPUT.IMAGE_SIZE,
cfg.INPUT.PIXEL_MEAN,
cfg.INPUT.PIXEL_STD)
target_transform = MatchPrior(
PriorBox(cfg)(), cfg.MODEL.CENTER_VARIANCE, cfg.MODEL.SIZE_VARIANCE,
cfg.MODEL.THRESHOLD)
if cfg.DATASETS.DG:
if args.eval_mode == "val":
dslist, val_set_dict = _create_dg_datasets(args, cfg, logger,
target_transform,
train_transform)
else:
dslist = _create_dg_datasets(args, cfg, logger, target_transform,
train_transform)
logger.info("Sizes of sources datasets:")
for k, v in dslist.items():
logger.info("{} size: {}".format(k, len(v)))
dataloaders = []
for name, train_dataset in dslist.items():
sampler = torch.utils.data.RandomSampler(train_dataset)
batch_sampler = torch.utils.data.sampler.BatchSampler(
sampler=sampler,
batch_size=cfg.SOLVER.BATCH_SIZE,
drop_last=True)
batch_sampler = samplers.IterationBasedBatchSampler(
batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER)
if cfg.MODEL.SELF_SUPERVISED:
ss_dataset = SelfSupervisedDataset(train_dataset, cfg)
train_loader = DataLoader(ss_dataset,
num_workers=args.num_workers,
batch_sampler=batch_sampler,
pin_memory=True)
else:
train_loader = DataLoader(train_dataset,
num_workers=args.num_workers,
batch_sampler=batch_sampler,
pin_memory=True)
dataloaders.append(train_loader)
if args.eval_mode == "val":
if args.return_best:
return do_train(cfg, model, dataloaders, optimizer, scheduler,
device, args, val_set_dict)
else:
return do_train(cfg, model, dataloaders, optimizer, scheduler,
device, args)
else:
return do_train(cfg, model, dataloaders, optimizer, scheduler,
device, args)
# No DG:
if args.eval_mode == "val":
train_dataset, val_dataset = build_dataset(
dataset_list=cfg.DATASETS.TRAIN,
transform=train_transform,
target_transform=target_transform,
split=True)
else:
train_dataset = build_dataset(dataset_list=cfg.DATASETS.TRAIN,
transform=train_transform,
target_transform=target_transform)
logger.info("Train dataset size: {}".format(len(train_dataset)))
if args.distributed:
sampler = torch.utils.data.DistributedSampler(train_dataset)
else:
sampler = torch.utils.data.RandomSampler(train_dataset)
batch_sampler = torch.utils.data.sampler.BatchSampler(
sampler=sampler, batch_size=cfg.SOLVER.BATCH_SIZE, drop_last=False)
batch_sampler = samplers.IterationBasedBatchSampler(
batch_sampler, num_iterations=cfg.SOLVER.MAX_ITER // args.num_gpus)
if cfg.MODEL.SELF_SUPERVISED:
ss_dataset = SelfSupervisedDataset(train_dataset, cfg)
train_loader = DataLoader(ss_dataset,
num_workers=args.num_workers,
batch_sampler=batch_sampler,
pin_memory=True)
else:
train_loader = DataLoader(train_dataset,
num_workers=args.num_workers,
batch_sampler=batch_sampler,
pin_memory=True)
if args.eval_mode == "val":
return do_train(cfg, model, train_loader, optimizer, scheduler, device,
args, {"validation_split": val_dataset})
else:
return do_train(cfg, model, train_loader, optimizer, scheduler, device,
args)
def run_demo(cfg, checkpoint_file, iou_threshold, score_threshold, images_dir,
output_dir):
basename = os.path.basename(checkpoint_file).split('.')[0]
epoch = basename[21:]
epoch = int(epoch)
if epoch < min_epoch:
return
device = torch.device(cfg.MODEL.DEVICE)
model = build_ssd_model(cfg)
checkpoint = torch.load(checkpoint_file)
model.load_state_dict(checkpoint['state_dict'])
print('Loaded weights from {}.'.format(checkpoint_file))
model = model.to(device)
model.eval()
predictor = Predictor(cfg=cfg,
model=model,
iou_threshold=iou_threshold,
score_threshold=score_threshold,
device=device)
cpu_device = torch.device("cpu")
image_paths = glob.glob(os.path.join(images_dir, '*.jpg'))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for image_path in tqdm(image_paths):
image = Image.open(image_path).convert("RGB")
image = np.array(image)
# image = image[:, ::-1]
#image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
output = predictor.predict(image)
boxes, scores, score_map = [o.to(cpu_device).numpy() for o in output]
# score_map = cv2.resize(score_map, (512, 512)) * 255
# score_map = score_map.astype(np.uint8)
# score_map = cv2.applyColorMap(score_map, cv2.COLORMAP_JET)
# score_map = cv2.resize(score_map,(1280,720),interpolation=cv2.INTER_CUBIC)
drawn_image = draw_bounding_boxes(image, boxes).astype(np.uint8)
image_name = os.path.basename(image_path)
txt_path = os.path.join(output_dir, 'txtes')
if not os.path.exists(txt_path):
os.makedirs(txt_path)
txt_path = os.path.join(txt_path,
'res_' + image_name.replace('jpg', 'txt'))
# print('txt_path:',txt_path)
with open(txt_path, 'w+') as f:
for box in boxes:
box_temp = np.reshape(box, (4, 2))
box = order_points_quadrangle(box_temp)
box = np.reshape(box, (1, 8)).squeeze(0)
is_valid = validate_clockwise_points(box)
if not is_valid:
continue
# print('box:',box)
line = ''
for item in box:
if item < 0:
item = 0
line += str(int(item)) + ','
line = line[:-1] + '\n'
f.write(line)
# path = os.path.join(output_dir, image_name)
# Image.fromarray(drawn_image).save(path)
# path = os.path.join(output_dir, image_name.split('.')[0]+'_score_map.'+image_name.split('.')[1])
# print(path)
# cv2.imwrite(path,score_map)
submit_path = MyZip(
'/home/binchengxiong/ssd_fcn_multitask_text_detection_pytorch1.0/demo/result223'
+ str(min_epoch) + '/txtes', epoch)
hmean_this_epoch = compute_hmean(submit_path, epoch)
