def detect_cv2(cfgfile, weightfile, imgfile, clsnamesfile=None, use_cuda=True):
import cv2
m = Darknet(cfgfile)
m.print_network()
m.is_yolov3 = True
m.load_weights(weightfile)
print('[INFO]:Loading weights from %s... Done!' % (weightfile))
if not clsnamesfile:
if m.num_classes == 20:
namesfile = 'data/voc.names'
elif m.num_classes == 80:
namesfile = 'data/coco.names'
else:
print(
'[INFO]:Default choose fail, You should give the class names file...'
)
return None
else:
namesfile = '%s' % clsnamesfile
if use_cuda:
m.cuda()
img = cv2.imread(imgfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
start = time.time()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
plot_boxes_cv2(img,
boxes,
savename='predictions.jpg',
class_names=class_names)
def detect(cfgfile,
weightfile,
imgfile,
num_classes,
clsnamesfile=None,
use_cuda=True):
m = Darknet(cfgfile)
m.print_network()
m.is_yolov3 = True
m.load_weights(weightfile)
print('[INFO]:Loading weights from %s... Done!' % (weightfile))
if not clsnamesfile:
if num_classes == 20:
namesfile = './data/voc.names'
elif num_classes == 80:
namesfile = './data/coco.names'
else:
print(
'[INFO]:Default choose fail, You should give the class names file...'
)
exit(-1)
else:
namesfile = '%s' % clsnamesfile
if use_cuda:
m.cuda()
img = Image.open(imgfile).convert('RGB')
sized = img.resize((m.width, m.height))
start = time.time()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
plot_boxes(img, boxes, 'predictions.jpg', class_names)
def detect(cfgfile, weightfile, imgfile, clsnamesfile=None, use_cuda=True):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('[INFO]:Loading weights from %s... Done!' % (weightfile))
if not clsnamesfile:
if m.num_classes == 20:
namesfile = 'data/voc.names'
elif m.num_classes == 80:
namesfile = 'data/coco.names'
else:
print(
'[INFO]:Default choose fail, You should give the class names file...'
)
exit(-1)
else:
namesfile = '%s' % clsnamesfile
if use_cuda:
m.cuda()
img = Image.open(imgfile).convert('RGB')
sized = img.resize((m.width, m.height))
conf_thresh = 0.2
nms_thresh = 0.3
m.eval()
if isinstance(sized, Image.Image):
sized = image2torch(sized)
elif type(sized) == np.ndarray:
sized = torch.from_numpy(sized.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
else:
print("[Error]: unknow image type...")
exit(-1)
if use_cuda:
sized = sized.cuda()
sized = torch.autograd.Variable(sized)
output = m(sized)
output = output.data
# by default, we only get one picture detected boxes.
anchors = [anchor / m.loss.stride for anchor in m.anchors]
boxes = get_region_boxes(output, conf_thresh, m.num_classes, anchors,
m.num_anchors)[0]
boxes = nms(boxes, nms_thresh)
class_names = load_class_names(namesfile)
plot_boxes(img, boxes, 'predictions.jpg', class_names)
def eval_list(cfgfile, weightfile, imglist, use_cuda=True):
m = Darknet(cfgfile)
m.eval()
m.load_weights(weightfile)
eval_wid = m.width
eval_hei = m.height
if use_cuda:
m.cuda()
conf_thresh = 0.25
nms_thresh = 0.4
iou_thresh = 0.5
min_box_scale = 8. / m.width
with open(imglist) as fp:
lines = fp.readlines()
total = 0.0
proposals = 0.0
correct = 0.0
lineId = 0
avg_iou = 0.0
for line in lines:
img_path = line.rstrip()
if img_path[0] == '#':
continue
lineId = lineId + 1
lab_path = img_path.replace('images', 'labels')
lab_path = lab_path.replace('JPEGImages', 'labels')
lab_path = lab_path.replace('.jpg', '.txt').replace('.png', '.txt')
truths = read_truths_args(lab_path, min_box_scale)
img = Image.open(img_path).convert('RGB').resize((eval_wid, eval_hei))
boxes = do_detect(m, img, conf_thresh, nms_thresh, use_cuda)
if False:
savename = "tmp/%06d.jpg" % (lineId)
print("save %s" % savename)
plot_boxes(img, boxes, savename)
total = total + truths.shape[0]
for i in range(len(boxes)):
if boxes[i][4] > conf_thresh:
proposals += 1
for i in range(truths.shape[0]):
box_gt = [truths[i][1], truths[i][2], truths[i][3], truths[i][4], 1.0]
best_iou = 0
for j in range(len(boxes)):
iou = bbox_iou(box_gt, boxes[j], x1y1x2y2=False)
best_iou = max(iou, best_iou)
if best_iou > iou_thresh:
avg_iou += best_iou
correct += 1
precision = 1.0*correct / proposals
recall = 1.0*correct / total
fscore = 2.0*precision*recall / (precision+recall)
print("%d IOU: %f, Recal: %f, Precision: %f, Fscore: %f\n" % (lineId-1, avg_iou/correct, recall, precision, fscore))
def Demo(cfgfile, weightfile, clsnamesfile=None, use_cuda=True):
m = Darknet(cfgfile)
# m.print_net()
m.load_weights(weightfile)
print('[INFO]:Loading weights from %s... Done!' % (weightfile))
if not clsnamesfile:
if m.num_classes == 20:
namesfile = 'data/voc.names'
elif m.num_classes == 80:
namesfile = 'data/coco.names'
else:
print('[INFO]:Default choose fail, You should give the class names file...')
exit(-1)
else:
namesfile = '%s' % clsnamesfile
class_names = load_class_names(namesfile)
if use_cuda:
m.cuda()
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('[Error]:Unable to open camera...')
exit(-1)
while True:
res, img = capture.read()
if res:
sized = cv2.resize(img, (m.width, m.height))
bboxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
draw_img = plot_boxes_cv2(img, bboxes, None, class_names)
cv2.imshow(cfgfile, draw_img)
cv2.waitKey(1)
import time
time.sleep(0.2)
else:
print('[Error]:Unable to read image...')
exit(-1)
class train():
def __init__(self, yolo_type='yolo2'):
self.yolo_type = yolo_type
if yolo_type == 'yolo1':
self.options = config.yolo1_options
elif yolo_type == 'yolo2':
self.options = config.yolo2_options
elif yolo_type == 'yolo3':
self.options = config.yolo3_options
else:
raise ValueError('Unknow yolo_type <%s>...' % yolo_type)
cfgparser = CfgParser()
self.net_options = cfgparser.parser(
cfgfile=self.options.get('cfgfile'))[0]
self.__initialization()
# start to train.
def start(self):
for epoch in range(self.init_epoch, self.max_epochs):
self.__train_epoch(epoch)
# train for one epoch.
def __train_epoch(self, epoch):
if self.use_cuda:
if self.ngpus > 1:
cur_model = self.model.module
else:
cur_model = self.model
else:
cur_model = self.model
train_loader = torch.utils.data.DataLoader(myDataset(
root=[self.trainSet, self.trainlabpth],
shape=(self.init_width, self.init_height),
shuffle=True,
transform=transforms.Compose([
transforms.ToTensor(),
]),
is_train=True,
seen=cur_model.seen,
num_workers=self.num_workers,
is_multiscale=self.is_multiscale,
jitter=self.jitter,
hue=self.hue,
saturation=self.saturation,
exposure=self.exposure,
max_object=self.max_object,
batch_size=self.batch_size),
batch_size=self.batch_size,
shuffle=False,
**self.kwargs)
logging(
'epoch %d, processed %d samples, processed_batches %d' %
(epoch, epoch * len(train_loader.dataset), self.processed_batches),
self.logsavefile)
self.model.train()
for batch_idx, (data, target) in enumerate(train_loader):
lr = self.__adjust_lr(self.optimizer, self.processed_batches)
self.processed_batches += 1
if self.use_cuda:
data = data.cuda()
data, target = Variable(data), Variable(target)
self.optimizer.zero_grad()
cur_model.seen += data.data.size(0)
loss = self.model(data, target)
if self.ngpus > 1:
loss = loss.sum()
loss.backward()
self.optimizer.step()
if ((epoch + 1) % self.save_interval == 0) or ((epoch + 1)
== self.max_epochs):
logging(
'save weights to %s/%06d.weights' %
(self.backupdir, epoch + 1), self.logsavefile)
cur_model.seen = (epoch + 1) * len(train_loader.dataset)
cur_model.save_weights('%s/%06d.weights' %
(self.backupdir, epoch + 1))
self.EM.eval(self.model)
# initialization
def __initialization(self):
self.use_cuda = self.options.get('use_cuda')
self.backupdir = self.options.get('backupdir')
self.gpus = self.options.get('gpus')
self.ngpus = self.options.get('ngpus')
self.save_interval = self.options.get('save_interval')
self.trainSet = self.options.get('trainSet')
self.testSet = self.options.get('testSet')
self.trainlabpth = self.options.get('trainlabpth')
self.testlabpth = self.options.get('testlabpth')
self.num_workers = self.options.get('num_workers')
self.is_multiscale = self.options.get('is_multiscale')
self.nsamples = file_lines(self.trainSet)
self.batch_size = int(self.net_options.get('batch'))
self.max_batches = int(self.net_options.get('max_batches'))
self.max_object = self.options.get('max_object')
self.learning_rate = float(self.net_options.get('learning_rate'))
self.steps = [
float(step) for step in self.net_options.get('steps').split(',')
]
self.scales = [
float(scale) for scale in self.net_options.get('scales').split(',')
]
self.momentum = float(self.net_options.get('momentum'))
self.decay = float(self.net_options.get('decay'))
self.jitter = self.options.get('jitter')
self.saturation = float(self.net_options.get('saturation'))
self.exposure = float(self.net_options.get('exposure'))
self.hue = float(self.net_options.get('hue'))
self.max_epochs = math.ceil(self.max_batches * self.batch_size /
self.nsamples)
self.logsavefile = self.options.get('logsavefile')
if not os.path.exists(self.backupdir):
os.mkdir(self.backupdir)
if self.use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = self.gpus
self.model = Darknet(self.options)
weightfile = self.options.get('weightfile')
if weightfile:
self.model.load_weights(weightfile)
print('[INFO]: %s loaded...' % weightfile)
self.init_width = int(self.net_options.get('width'))
self.init_height = int(self.net_options.get('height'))
self.init_epoch = self.model.seen // self.nsamples
self.processed_batches = self.model.seen // self.batch_size
self.kwargs = {
'num_workers': self.num_workers,
'pin_memory': True
} if self.use_cuda else {}
if self.use_cuda:
if self.ngpus > 1:
self.model = nn.DataParallel(self.model).cuda()
else:
self.model = self.model.cuda()
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.learning_rate / self.batch_size,
momentum=self.momentum,
dampening=0,
weight_decay=self.decay * self.batch_size)
self.EM = evalModel(num_workers=self.num_workers,
use_cuda=self.use_cuda,
testSet=self.testSet,
testlabpth=self.testlabpth,
init_width=self.init_width,
init_height=self.init_height,
batch_size=self.batch_size,
ngpus=self.ngpus,
by_stride=self.options.get('by_stride'),
conf_thresh=self.options.get('conf_thresh'),
iou_thresh=self.options.get('iou_thresh'),
max_object=self.max_object,
is_multiscale=self.is_multiscale,
yolo_type=self.yolo_type,
nms_thresh=self.options.get('nms_thresh'),
logsavefile=self.logsavefile)
# adjust learning rate.
def __adjust_lr(self, optimizer, batch_idx):
lr = self.learning_rate
for i in range(len(self.steps)):
scale = self.scales[i] if i < len(self.scales) else 1
if batch_idx >= self.steps[i]:
lr = lr * scale
if batch_idx == self.steps[i]:
break
else:
break
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr / self.batch_size
return lr
def __initialization(self):
self.use_cuda = self.options.get('use_cuda')
self.backupdir = self.options.get('backupdir')
self.gpus = self.options.get('gpus')
self.ngpus = self.options.get('ngpus')
self.save_interval = self.options.get('save_interval')
self.trainSet = self.options.get('trainSet')
self.testSet = self.options.get('testSet')
self.trainlabpth = self.options.get('trainlabpth')
self.testlabpth = self.options.get('testlabpth')
self.num_workers = self.options.get('num_workers')
self.is_multiscale = self.options.get('is_multiscale')
self.nsamples = file_lines(self.trainSet)
self.batch_size = int(self.net_options.get('batch'))
self.max_batches = int(self.net_options.get('max_batches'))
self.max_object = self.options.get('max_object')
self.learning_rate = float(self.net_options.get('learning_rate'))
self.steps = [
float(step) for step in self.net_options.get('steps').split(',')
]
self.scales = [
float(scale) for scale in self.net_options.get('scales').split(',')
]
self.momentum = float(self.net_options.get('momentum'))
self.decay = float(self.net_options.get('decay'))
self.jitter = self.options.get('jitter')
self.saturation = float(self.net_options.get('saturation'))
self.exposure = float(self.net_options.get('exposure'))
self.hue = float(self.net_options.get('hue'))
self.max_epochs = math.ceil(self.max_batches * self.batch_size /
self.nsamples)
self.logsavefile = self.options.get('logsavefile')
if not os.path.exists(self.backupdir):
os.mkdir(self.backupdir)
if self.use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = self.gpus
self.model = Darknet(self.options)
weightfile = self.options.get('weightfile')
if weightfile:
self.model.load_weights(weightfile)
print('[INFO]: %s loaded...' % weightfile)
self.init_width = int(self.net_options.get('width'))
self.init_height = int(self.net_options.get('height'))
self.init_epoch = self.model.seen // self.nsamples
self.processed_batches = self.model.seen // self.batch_size
self.kwargs = {
'num_workers': self.num_workers,
'pin_memory': True
} if self.use_cuda else {}
if self.use_cuda:
if self.ngpus > 1:
self.model = nn.DataParallel(self.model).cuda()
else:
self.model = self.model.cuda()
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.learning_rate / self.batch_size,
momentum=self.momentum,
dampening=0,
weight_decay=self.decay * self.batch_size)
self.EM = evalModel(num_workers=self.num_workers,
use_cuda=self.use_cuda,
testSet=self.testSet,
testlabpth=self.testlabpth,
init_width=self.init_width,
init_height=self.init_height,
batch_size=self.batch_size,
ngpus=self.ngpus,
by_stride=self.options.get('by_stride'),
conf_thresh=self.options.get('conf_thresh'),
iou_thresh=self.options.get('iou_thresh'),
max_object=self.max_object,
is_multiscale=self.is_multiscale,
yolo_type=self.yolo_type,
nms_thresh=self.options.get('nms_thresh'),
logsavefile=self.logsavefile)
def valid(datacfg, cfgfile, weightfile, outfile, use_cuda=True):
options = read_data_cfg(datacfg)
valid_images = options['valid']
name_list = options['names']
prefix = 'results'
names = load_class_names(name_list)
with open(valid_images) as fp:
tmp_files = fp.readlines()
valid_files = [item.rstrip() for item in tmp_files]
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
if use_cuda:
m.cuda()
m.eval()
valid_dataset = dataset.listDataset(valid_images,
shape=(m.width, m.height),
shuffle=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
valid_batchsize = 2
assert (valid_batchsize > 1)
kwargs = {'num_workers': 4, 'pin_memory': True}
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=valid_batchsize,
shuffle=False,
**kwargs)
fps = [0] * m.num_classes
if not os.path.exists(prefix):
os.mkdir(prefix)
for i in range(m.num_classes):
buf = '%s/%s%s.txt' % (prefix, outfile, names[i])
fps[i] = open(buf, 'w')
lineId = -1
conf_thresh = 0.005
nms_thresh = 0.45
for batch_idx, (data, target) in enumerate(valid_loader):
if use_cuda:
data = data.cuda()
data = Variable(data, volatile=True)
output = m(data).data
batch_boxes = get_region_boxes(output, conf_thresh, m.num_classes,
m.anchors, m.num_anchors, 0, 1)
for i in range(output.size(0)):
lineId = lineId + 1
fileId = os.path.basename(valid_files[lineId]).split('.')[0]
width, height = get_image_size(valid_files[lineId])
print(valid_files[lineId])
boxes = batch_boxes[i]
boxes = nms(boxes, nms_thresh)
for box in boxes:
x1 = (box[0] - box[2] / 2.0) * width
y1 = (box[1] - box[3] / 2.0) * height
x2 = (box[0] + box[2] / 2.0) * width
y2 = (box[1] + box[3] / 2.0) * height
det_conf = box[4]
for j in range((len(box) - 5) // 2):
cls_conf = box[5 + 2 * j]
cls_id = box[6 + 2 * j]
prob = det_conf * cls_conf
fps[cls_id].write('%s %f %f %f %f %f\n' %
(fileId, prob, x1, y1, x2, y2))
for i in range(m.num_classes):
fps[i].close()
# Draw network structure
from graphviz import Digraph
import torch
from torch.autograd import Variable
from torchviz import make_dot
from nets.darknet import Darknet
model = Darknet('./cfg/me/darknet19_wfc_face.cfg')
model.print_network()
x = Variable(torch.randn(1, 3, 416, 416), requires_grad=True)
y = model(x)
params_dict = dict(model.named_parameters())
params_dict['x'] = x
g = make_dot(y, params=params_dict)
g.view()
def Demo(yolo_type='yolo2'):
# for yoloV1
if yolo_type == 'yolo1':
options = config.yolo1_options
weightfile = options.get('weightfile')
clsnamesfile = options.get('clsnamesfile')
use_cuda = options.get('use_cuda')
conf_thresh = options.get('conf_thresh')
nms_thresh = options.get('nms_thresh')
model = Darknet(options)
width = int(model.blocks[0]['width'])
height = int(model.blocks[0]['height'])
num_classes = int(model.blocks[-1]['classes'])
num_anchors = int(model.blocks[-1]['num'])
if (not weightfile) or (not clsnamesfile):
print(
'[Error]: You should assign the weightfile and clsnamesfile in config.py-yolo1_options...'
)
sys.exit(-1)
model.load_weights(options.get('weightfile'))
print('[INFO]: Loading weights from %s... Done!' % weightfile)
class_names = load_class_names(clsnamesfile)
if use_cuda:
model.cuda()
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('[Error]:Unable to open camera...')
sys.exit(-1)
while True:
res, img = capture.read()
if res:
sized = cv2.resize(img, (width, height))
model.eval()
if isinstance(sized, Image.Image):
sized = image2torch(sized)
elif type(sized) == np.ndarray:
sized = torch.from_numpy(sized.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
else:
print("[Error]: unknow image type...")
sys.exit(-1)
if use_cuda:
sized = sized.cuda()
sized = torch.autograd.Variable(sized)
output = model(sized)
output = output.data
boxes = get_boxes_yolo1(output,
conf_thresh=conf_thresh,
num_classes=num_classes,
num_anchors=num_anchors,
width=width,
height=height,
stride=model.det_strides[0])[0]
bboxes = nms(boxes, nms_thresh)
draw_img = plot_boxes_cv2(img, bboxes, class_names)
cv2.imshow('yolo1', draw_img)
cv2.waitKey(1)
else:
sys.exit(-1)
# for yoloV2
elif yolo_type == 'yolo2':
options = config.yolo2_options
weightfile = options.get('weightfile')
clsnamesfile = options.get('clsnamesfile')
use_cuda = options.get('use_cuda')
conf_thresh = options.get('conf_thresh')
nms_thresh = options.get('nms_thresh')
by_stride = options.get('by_stride')
model = Darknet(options)
width = int(model.blocks[0]['width'])
height = int(model.blocks[0]['height'])
anchors = [float(i) for i in model.blocks[-1]['anchors'].split(',')]
num_classes = int(model.blocks[-1]['classes'])
num_anchors = int(model.blocks[-1]['num'])
if (not weightfile) or (not clsnamesfile):
print(
'[Error]: You should assign the weightfile and clsnamesfile in config.py-yolo2_options...'
)
sys.exit(-1)
model.load_weights(options.get('weightfile'))
print('[INFO]: Loading weights from %s... Done!' % weightfile)
class_names = load_class_names(clsnamesfile)
if use_cuda:
model.cuda()
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('[Error]:Unable to open camera...')
sys.exit(-1)
while True:
res, img = capture.read()
if res:
sized = cv2.resize(img, (width, height))
model.eval()
if isinstance(sized, Image.Image):
sized = image2torch(sized)
elif type(sized) == np.ndarray:
sized = torch.from_numpy(sized.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
else:
print("[Error]: unknow image type...")
sys.exit(-1)
if use_cuda:
sized = sized.cuda()
sized = torch.autograd.Variable(sized)
output = model(sized)
output = output.data
boxes = get_boxes_yolo2(output,
conf_thresh=conf_thresh,
num_classes=num_classes,
anchors=anchors,
num_anchors=num_anchors,
stride=model.det_strides[0],
by_stride=by_stride)[0]
bboxes = nms(boxes, nms_thresh)
draw_img = plot_boxes_cv2(img, bboxes, class_names)
cv2.imshow('yolo2', draw_img)
cv2.waitKey(1)
else:
sys.exit(-1)
# for yoloV3
elif yolo_type == 'yolo3':
options = config.yolo3_options
weightfile = options.get('weightfile')
clsnamesfile = options.get('clsnamesfile')
use_cuda = options.get('use_cuda')
conf_thresh = options.get('conf_thresh')
nms_thresh = options.get('nms_thresh')
by_stride = options.get('by_stride')
model = Darknet(options)
width = int(model.blocks[0]['width'])
height = int(model.blocks[0]['height'])
anchors = [float(i) for i in model.blocks[-1]['anchors'].split(',')]
anchor_masks = []
for block in model.blocks:
if block['layer_type'] == 'yolo':
anchor_mask = [int(i) for i in block['mask'].split(',')]
anchor_masks.append(anchor_mask)
num_classes = int(model.blocks[-1]['classes'])
num_anchors = int(model.blocks[-1]['num'])
if (not weightfile) or (not clsnamesfile):
print(
'[Error]: You should assign the weightfile and clsnamesfile in config.py-yolo3_options...'
)
sys.exit(-1)
model.load_weights(options.get('weightfile'))
print('[INFO]: Loading weights from %s... Done!' % weightfile)
class_names = load_class_names(clsnamesfile)
if use_cuda:
model.cuda()
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('[Error]:Unable to open camera...')
sys.exit(-1)
while True:
res, img = capture.read()
if res:
sized = cv2.resize(img, (width, height))
model.eval()
if isinstance(sized, Image.Image):
sized = image2torch(sized)
elif type(sized) == np.ndarray:
sized = torch.from_numpy(sized.transpose(
2, 0, 1)).float().div(255.0).unsqueeze(0)
else:
print("[Error]: unknow image type...")
sys.exit(-1)
if use_cuda:
sized = sized.cuda()
sized = torch.autograd.Variable(sized)
output = model(sized)
output = output
boxes = get_boxes_yolo3(output,
conf_thresh=conf_thresh,
num_classes=num_classes,
anchors=anchors,
num_anchors=num_anchors,
anchor_masks=anchor_masks,
stride=model.det_strides,
by_stride=by_stride)
bboxes = []
for bs in boxes:
bboxes += bs
bboxes = nms(bboxes, nms_thresh)
draw_img = plot_boxes_cv2(img, bboxes, class_names)
cv2.imshow('yolo3', draw_img)
cv2.waitKey(1)
else:
sys.exit(-1)
else:
raise ValueError('yolo_type should be in [yolo1, yolo2, yolo3]...')
seed = Settings.seed
eps = Settings.eps
# test
conf_thresh = Settings.conf_thresh
nms_thresh = Settings.nms_thresh
iou_thresh = Settings.iou_thresh
# --------------------------------------------------------------------------------------------------------------------
# some pre-process
# --------------------------------------------------------------------------------------------------------------------
torch.manual_seed(seed)
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed)
model = Darknet(cfgfile)
model.print_network()
model.is_yolov3 = YOLOv3
init_width = model.width
init_height = model.height
kwargs = {'num_workers': num_workers, 'pin_memory': True} if use_cuda else {}
test_loader = torch.utils.data.DataLoader(
dataset.listDataset(testlist,
shape=(init_width, init_height),
shuffle=False,
transform=transforms.Compose([transforms.ToTensor(),]),
train=False),
batch_size=batch_size,
shuffle=False,
**kwargs)
if use_cuda:
# Test HParams
conf_thresh = Settings.conf_thresh
nms_thresh = Settings.nms_thresh
iou_thresh = Settings.iou_thresh
# --------------------------------------------------------------------------------------------------------------------
# Some pre-process
# --------------------------------------------------------------------------------------------------------------------
if not os.path.exists(backupdir):
os.mkdir(backupdir)
# Sets the seed for generating random numbers. And returns a torch._C.Generator object.
torch.manual_seed(seed)
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed)
model = Darknet(cfgfile)
model.print_network()
# Different yolo version
if YOLOv3:
yolo_loss = model.losses
else:
region_loss = model.loss
try:
model.load_weights(weightfile)
print('[INFO]:%s loaded...' % weightfile)
except:
pass
model.is_yolov3 = YOLOv3
if YOLOv3:
for yl in yolo_loss:
yl.seen = model.seen
img = self.transform(img)
if self.target_transform is not None:
label = self.target_transform(label)
if self.num_workers > 0:
self.seen = self.seen + self.num_workers
elif self.num_workers == 0:
self.seen += 1
else:
print('[Error]:num_workers must greater than zero...')
return (img, label)
if __name__ == '__main__':
batch_size = 1
testlist = '2007_test.txt'
init_width = 416
init_height = 416
from torchvision import transforms
from nets.darknet import Darknet
cfgfile = './cfg/yolov3.cfg'
num_workers = 4
model = Darknet(cfgfile)
cur_model = model
kwargs = {'num_workers': num_workers, 'pin_memory': True}
listDataset(testlist,
shape=(init_width, init_height),
shuffle=False,
transform=transforms.Compose([
transforms.ToTensor(),
]),
train=False)[0]