def run(img_dir, output_dir, img_size, num_classes, weights, conf_thres,
nms_thres, show):
shutil.rmtree(output_dir, ignore_errors=True)
os.makedirs(output_dir, exist_ok=True)
model = YOLOV3(num_classes, img_size)
state_dict = torch.load(weights, map_location='cpu')
model.load_state_dict(state_dict['model'])
model = model.to(device)
model.eval()
colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(80)]
names = [n for n in os.listdir(img_dir) if osp.splitext(n)[1] in IMG_EXT]
names.sort()
for name in tqdm(names):
img = cv2.imread(osp.join(img_dir, name))
det = inference(model, [img], img_size, conf_thres, nms_thres)[0]
det_txt = []
# Write results
for *xyxy, conf, _, cls in det:
det_txt.append(' '.join(['%g'] * 6) % (*xyxy, cls, conf))
if show: # Add bbox to image
label = '%d %.2f' % (int(cls), conf)
plot_one_box(xyxy, img, label=label, color=colors[int(cls)])
with open(osp.join(output_dir,
osp.splitext(name)[0] + '.txt'), 'w') as f:
f.write('\n'.join(det_txt))
# Stream results
if show:
cv2.imshow('yolo', img)
cv2.waitKey(1)
# Save results (image with detections)
cv2.imwrite(osp.join(output_dir, name), img)
def export2caffe(weights, num_classes, img_size):
os.environ['MODEL_EXPORT'] = '1'
model = YOLOV3(num_classes)
weights = torch.load(weights, map_location='cpu')
model.load_state_dict(weights['model'])
model.eval()
fuse(model)
name = 'RYOLOV3'
dummy_input = torch.ones([1, 3, img_size[1], img_size[0]])
pytorch2caffe.trans_net(model, dummy_input, name)
pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name))
def train(data_dir, epochs, img_size, batch_size, accumulate, lr, adam, resume,
weights, num_workers, multi_scale, rect, mixed_precision, notest,
nosave):
train_coco = osp.join(data_dir, 'train.json')
val_coco = osp.join(data_dir, 'val.json')
train_data = CocoDataset(train_coco,
img_size=img_size,
multi_scale=multi_scale,
rect=rect,
with_label=True,
mosaic=False)
train_loader = DataLoader(
train_data,
batch_size=batch_size,
shuffle=not (dist.is_initialized()),
sampler=DistributedSampler(train_data, dist.get_world_size(),
dist.get_rank())
if dist.is_initialized() else None,
pin_memory=True,
num_workers=num_workers,
collate_fn=CocoDataset.collate_fn,
)
train_fetcher = Fetcher(train_loader, train_data.post_fetch_fn)
if not notest:
val_data = CocoDataset(val_coco,
img_size=img_size,
augments=None,
rect=rect)
val_loader = DataLoader(
val_data,
batch_size=batch_size,
shuffle=not (dist.is_initialized()),
sampler=DistributedSampler(val_data, dist.get_world_size(),
dist.get_rank())
if dist.is_initialized() else None,
pin_memory=True,
num_workers=num_workers,
collate_fn=CocoDataset.collate_fn,
)
val_fetcher = Fetcher(val_loader, post_fetch_fn=val_data.post_fetch_fn)
model = YOLOV3(len(train_data.classes))
trainer = Trainer(model,
train_fetcher,
loss_fn=compute_loss,
workdir='weights',
accumulate=accumulate,
adam=adam,
lr=lr,
weights=weights,
resume=resume,
mixed_precision=mixed_precision)
while trainer.epoch < epochs:
trainer.step()
best = False
if not notest:
metrics = test(trainer.model, val_fetcher, conf_thres=0.1)
if metrics > trainer.metrics:
best = True
print('save best, mAP: %g' % metrics)
trainer.metrics = metrics
if not nosave:
trainer.save(best)
type=float,
default=0.1,
help='object confidence threshold')
parser.add_argument('--nms-thres',
type=float,
default=0.5,
help='iou threshold for non-maximum suppression')
opt = parser.parse_args()
val_data = CocoDataset(opt.val,
img_size=opt.img_size,
augments=None,
rect=opt.rect)
val_loader = DataLoader(
val_data,
batch_size=opt.batch_size,
pin_memory=True,
num_workers=opt.num_workers,
collate_fn=CocoDataset.collate_fn,
)
val_fetcher = Fetcher(val_loader, post_fetch_fn=val_data.post_fetch_fn)
model = YOLOV3(len(val_data.classes))
if opt.weights:
state_dict = torch.load(opt.weights, map_location='cpu')
model.load_state_dict(state_dict['model'])
metrics = test(model,
val_fetcher,
conf_thres=opt.conf_thres,
nms_thres=opt.nms_thres)
print('metrics: %8g' % (metrics))