def main():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', default='', type=str)
parser.add_argument('--weight_dir', default='weights', type=str)
parser.add_argument('--test_dir', default='test', type=str)
parser.add_argument('--data_dir', default='data', type=str)
parser.add_argument('--gpu', default='', type=str)
parser.add_argument('--image',
default='',
type=str,
help='path to test image')
parser.add_argument('--video',
default='',
type=str,
help='path to test video')
args = parser.parse_args()
if args.gpu is not '':
cfg.GPU = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU
net = DarkNet(False)
weight_file = os.path.join(args.data_dir, args.weight_dir, args.weights)
test_path = os.path.join(args.data_dir, args.test_dir)
detector = Detector(net, weight_file)
if args.image is not '':
imname = os.path.join(test_path, args.image)
detector.image_detector(imname)
if args.video is not '':
video = os.path.join(test_path, args.video)
cap = cv2.VideoCapture(video)
detector.video_detector(cap)
def __init__(self,
model_path,
class_name_list=None,
mean_rgb=[122.67891434, 116.66876762, 104.00698793],
conf_thresh=0.1,
prob_thresh=0.1,
nms_thresh=0.5,
gpu_id=0):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
use_gpu = torch.cuda.is_available()
assert use_gpu, 'Current implementation does not support CPU mode. Enable CUDA.'
# Load YOLO model.
print("Loading YOLO model...")
darknet = DarkNet(conv_only=True, bn=True, init_weight=True)
darknet.features = torch.nn.DataParallel(darknet.features)
self.yolo = YOLOv1(darknet.features)
self.yolo.conv_layers = torch.nn.DataParallel(self.yolo.conv_layers)
self.yolo.load_state_dict(torch.load(model_path))
self.yolo.cuda()
print("Done loading!")
self.yolo.eval()
self.S = self.yolo.feature_size
self.B = self.yolo.num_bboxes
self.C = self.yolo.num_classes
self.class_name_list = class_name_list if (
class_name_list is not None) else list(VOC_CLASS_BGR.keys())
assert len(self.class_name_list) == self.C
self.mean = np.array(mean_rgb, dtype=np.float32)
assert self.mean.shape == (3, )
self.conf_thresh = conf_thresh
self.prob_thresh = prob_thresh
self.nms_thresh = nms_thresh
self.to_tensor = transforms.ToTensor()
# Warm up.
dummy_input = Variable(torch.zeros((1, 3, 448, 448)))
dummy_input = dummy_input.cuda()
for i in range(10):
self.yolo(dummy_input)
def main():
args = parser.parse_args()
use_cuda = torch.cuda.is_available() and not args.no_cuda
model = DarkNet(use_cuda)
model_source = torch.load(args.weights)
model.load_state_dict(model_source['model'])
model.eval()
if use_cuda:
model = model.cuda()
ip = IMGProcess(model_source,
use_cuda=use_cuda,
img_path=args.images,
img_size=args.img_size,
confidence=args.confidence,
result=args.result)
print("-" * 57 + "Result" + "-" * 57)
for batch in ip:
outputs = ip.predict(model(batch), nms_conf=args.nms_thresh)
for name, objs in outputs:
print("Image - {}".format(name))
print("Detect Objects - [{}]".format(", ".join(objs)))
print("-" * 120)
def main():
classes = load_coco_names("data/coco.names")
tf.disable_v2_behavior()
inputs = tf.placeholder(
tf.float32, [None, _IMG_DIM, _IMG_DIM, 3], name="inputs")
with tf.variable_scope('detector'):
# build model
darknet = DarkNet("cfg/yolov3.cfg")
darknet.build(inputs)
load_ops = load_weights(tf.global_variables(
scope='detector'), 'weights/yolov3.weights')
saver = tf.train.Saver(tf.global_variables(scope='detector'))
with tf.Session() as sess:
sess.run(load_ops)
save_path = saver.save(sess, save_path='./weights/model.ckpt')
def detect(cfgfile, weightfile, imgfile):
m = DarkNet(cfgfile)
#m.print_net()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if m.num_classes == 20:
namesfile = '../data/voc.names'
elif m.num_classes == 80:
namesfile = '../data/coco.names'
else:
namesfile = '../data/names'
img = Image.open(imgfile).convert('RGB')
sized = letterbox_image(img, m.width, m.height)
start = time.time()
use_cuda = torch.cuda.is_available()
if use_cuda:
m.cuda()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
correct_yolo_boxes(boxes, img.width, img.height, m.width, m.height)
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)
predictions_img = Image.open('predictions.jpg')
predictions_img.show()
def detect(cfgfile, weightfile, imgfile):
"""
- weight file was downloaded at https://pjreddie.com/darknet/yolov2/
:cfgfile network config file
:weightfile weight file, include w and b
:imgfile single image input
"""
# init model
model = DarkNet(cfgfile)
model.print_network()
model.load_weights(weightfile)
print 'load weights done!'
if model.num_class == 20:
namesfile = 'cfg/voc.names'
elif model.num_class == 80:
namesfile = 'cfg/coco.names'
img = Image.open(imgfile).convert('RGB')
# resize to 416x416
img_sized = img.resize((model.width, model.height))
# get model output
output = get_net_output(model, img_sized)
# get region boxes
boxes = get_region_boxes(output, 0.5, model.num_class, model.anchors,
model.num_anchor)[0]
# iou > 0.4
boxes = nms(boxes, 0.4)
for i in range(len(boxes)):
print boxes[i]
class_names = load_class_names(namesfile)
# plot bounding box
plot_boxes(img, boxes, 'predictions.jpg', class_names)
# show img
predictions_img = Image.open('predictions.jpg')
predictions_img.show()
def test():
from torch.autograd import Variable
# Build model with randomly initialized weights
darknet = DarkNet(conv_only=True, bn=True, init_weight=True)
yolo = YOLOv1(darknet.features)
# Prepare a dummy image to input
image = torch.rand(10, 3, 448, 448)
image = Variable(image)
# Forward
output = yolo(image)
# Check ouput tensor size, which should be [10, 7, 7, 30]
print(output.size())
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', default='', type=str)
parser.add_argument('--data_dir', default='data', type=str)
parser.add_argument('--gpu', default='', type=str)
args = parser.parse_args()
if args.data_dir != cfg.DATA_PATH:
update_cfgpath(args.data_dir, args.weights)
if args.gpu is not '':
cfg.GPU = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU
solver = Solver(DarkNet(), pascal_voc('train'))
print('==== Start Training ====')
solver.train()
print('==== Finish Training ====')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', default='', type=str)
parser.add_argument('--weight_dir', default='weights', type=str)
parser.add_argument('--data_dir', default='data', type=str)
parser.add_argument('--gpu', default='', type=str)
args = parser.parse_args()
if args.gpu is not '':
cfg.GPU = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = cfg.GPU
darknet = DarkNet(False)
weight_file = os.path.join(args.data_dir, args.weight_dir, args.weights)
data = pascal_voc('test')
evaluater = Evaluater(darknet, weight_file, data)
print('==== Start evaluation ====')
evaluater.test()
print('==== Finish evaluation ====')
"""
def camera_detect(cfgfile, weightfile):
"""
- camera detect
:cfgfile use tiny config file
:weightfile use tiny weight file
"""
model = DarkNet(cfgfile)
model.print_net()
model.load_weights(weightfile)
print('load weights done!')
num_classes = 80
if num_classes == 20:
namesfile = '../data/voc.names'
elif num_classes == 80:
namesfile = '../data/coco.names'
else:
namesfile = '../data/names'
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("Unable to open camera")
exit(-1)
while True:
res, img = cap.read()
if res:
img = Image.fromarray(img, mode='RGB') # numpy.array -> PIL.Image
sized = letterbox_image(img, model.width, model.height)
boxes = do_detect(model, sized, 0.5, 0.4, False)
correct_yolo_boxes(boxes, img.width, img.height, model.width,
model.height)
class_names = load_class_names(namesfile)
image_draw = plot_boxes(img, boxes, None, class_names)
np_img = np.asarray(image_draw) # PIL.Image -> numpy.array
cv2.imshow(cfgfile, np_img)
cv2.waitKey(1)
else:
print("Unable to read image")
exit(-1)
def camera_detect(cfgfile, weightfile):
"""
- camera detect
:cfgfile use tiny config file
:weightfile use tiny weight file
"""
model = DarkNet(cfgfile)
model.print_network()
model.load_weights(weightfile)
print 'load weights done!'
if model.num_class == 20:
namesfile = 'cfg/voc.names'
elif model.num_class == 80:
namesfile = 'cfg/coco.names'
class_names = load_class_names(namesfile)
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("Unable to open camera")
exit(-1)
while True:
res, img = cap.read()
if res:
img_sized = cv2.resize(img, (model.width, model.height))
# get output
output = get_net_output(model, img_sized)
boxes = get_region_boxes(output, 0.5, model.num_class,
model.anchors, model.num_anchor)[0]
print 'boxes:', boxes
boxes = nms(boxes, 0.4)
print('-------------')
draw_img = plot_boxes_cv2(img, boxes, None, class_names)
cv2.imshow(cfgfile, draw_img)
cv2.waitKey(1)
else:
print("Unable to read image")
exit(-1)
# Test parameters
conf_thresh = 0.25
nms_thresh = 0.4
iou_thresh = 0.5
if not os.path.exists(backupdir):
os.mkdir(backupdir)
###############
torch.manual_seed(seed)
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed)
model = DarkNet(cfgfile)
region_loss = model.loss
model.load_weights(weightfile)
model.print_network()
region_loss.seen = model.seen
processed_batches = model.seen / batch_size
init_width = model.width
init_height = model.height
init_epoch = model.seen / nsamples
kwargs = {'num_workers': num_workers, 'pin_memory': True} if use_cuda else {}
test_loader = torch.utils.data.DataLoader(dataset.ListDataset(
testlist,
if __name__ == "__main__":
args = arg_parse()
images = args.images
batch_size = int(args.batch_size)
confidence = float(args.confidence)
nms_thresh = float(args.nms_thresh)
start = 0
CUDA = torch.cuda.is_available()
num_classes = 80 #For COCO
classes = load_classes("../data/coco.names")
classes[0] = 'beautiful girl'
# 获取network 加载参数
print("Loading network......")
model = DarkNet(args.cfgfile)
model.load_weights(args.weightsfile)
print('Nerwork successfully loaded')
model.net_info['height'] = args.reso
inp_dim = int(model.net_info['height'])
assert inp_dim % 32 == 0
assert inp_dim > 32
if CUDA:
model = model.cuda()
model.eval() # 梯度不计算， drop out #
# 读取图片
read_dir = time.time()
try:
imlist = [
osp.join(osp.realpath('.'), images, img)
return img
if __name__ == '__main__':
args = arg_parse()
images = args.images
bs = args.bs
confidence = args.confidence
nms_thresh = args.nms_thresh
cuda = torch.cuda.is_available()
classes = load_classes('./cfg/coco.names')
num_classes = 80
#set up the network
print('loading the network')
model = DarkNet(args.cfg_file)
model.load_weights(args.weight)
print('load the network successfully')
model.net_info['height'] = args.reso
inp_dim = int(model.net_info['height'])
assert inp_dim % 32 == 0
assert inp_dim > 32
if cuda:
model.cuda()
#set the model into eval mode
model.eval()
read_dir = time.time()
def train(folder="weights"):
os.makedirs(folder, exist_ok=True)
args = parser.parse_args()
use_cuda = torch.cuda.is_available() and not args.no_cuda
classes = load_classes()
num_classes = len(classes)
model = DarkNet(use_cuda, num_classes)
if use_cuda:
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters())
training_data = Data_loader("data/labels/train2014/",
"data/train2014",
img_size=args.img_size,
max_objects=args.max_objects,
batch_size=args.batch_size,
is_cuda=use_cuda)
validation_data = Data_loader("data/labels/val2014/",
"data/val2014",
img_size=args.img_size,
max_objects=args.max_objects,
batch_size=args.batch_size,
is_cuda=use_cuda)
for epoch in range(args.epoch):
model.train()
for batch_i, (imgs, labels) in enumerate(training_data):
optimizer.zero_grad()
loss, gather_losses = model(imgs, labels)
loss.backward()
optimizer.step()
print(
f"""[Epoch {epoch+1}/{args.epoch},Batch {batch_i+1}/{training_data.stop_step}] [Losses: x {gather_losses["x"]:.5f}, y {gather_losses["y"]:.5f}, w {gather_losses["w"]:.5f}, h { gather_losses["h"]:.5f}, conf {gather_losses["conf"]:.5f}, cls {gather_losses["cls"]:.5f}, total {loss.item():.5f}, recall: {gather_losses["recall"]:.5f}, precision: {gather_losses["precision"]:.5f}]"""
)
torch.save({
"model": model.state_dict(),
"classes": classes
}, f"{folder}/{epoch}.weights.pt")
all_detections = []
all_annotations = []
model.eval()
for imgs, labels in validation_data:
with torch.no_grad():
prediction = model(imgs)
outputs = predict(prediction, args.nms_conf, args.confidence)
labels = labels.cpu()
for output, annotations in zip(outputs, labels):
all_detections.append(
[np.array([]) for _ in range(num_classes)])
if output is not None:
pred_boxes = output[:, :5].cpu().numpy()
scores = output[:, 4].cpu().numpy()
pred_labels = output[:, -1].cpu().numpy()
sort_i = np.argsort(scores)
pred_labels = pred_labels[sort_i]
pred_boxes = pred_boxes[sort_i]
for label in range(num_classes):
all_detections[-1][label] = pred_boxes[pred_labels ==
label]
all_annotations.append(
[np.array([]) for _ in range(num_classes)])
if any(annotations[:, -1] > 0):
annotation_labels = annotations[annotations[:, -1] > 0,
0].numpy()
_annotation_boxes = annotations[annotations[:, -1] > 0, 1:]
annotation_boxes = np.empty_like(_annotation_boxes)
annotation_boxes[:,
0] = _annotation_boxes[:,
0] - _annotation_boxes[:,
2] / 2
annotation_boxes[:,
1] = _annotation_boxes[:,
1] - _annotation_boxes[:,
3] / 2
annotation_boxes[:,
2] = _annotation_boxes[:,
0] + _annotation_boxes[:,
2] / 2
annotation_boxes[:,
3] = _annotation_boxes[:,
1] + _annotation_boxes[:,
3] / 2
annotation_boxes *= args.img_size
for label in range(num_classes):
all_annotations[-1][label] = annotation_boxes[
annotation_labels == label, :]
average_precisions = evaluate(num_classes, all_detections,
all_annotations)
print(f"""{"-"*40}evaluation.{epoch}{"-"*40}""")
for c, ap in average_precisions.items():
print(f"Class '{c}' - AP: {ap}")
mAP = np.mean(list(average_precisions.values()))
print(f"mAP: {mAP}")
print(f"""{"-"*40}end{"-"*40}""")
confidence = float(args.confidence)
# NMS阀值
nms_thesh = float(args.nms_thresh)
start = 0
# 是否有可用的GPU
CUDA = torch.cuda.is_available()
# 检测物标种类
num_classes = 80
# 导入coco的80个物标种类名字
classes = load_classes('data/coco.names')
# 创建神经网络
print("Loading network.....")
# 备用配置文件
model = DarkNet(args.cfgfile)
# 导入DarkNet网络权重参数
model.load_weights(args.weightsfile)
print("Network successfully loaded")
# 图片的分辨率，默认为416x416.一般为32的z倍数
# 输入图像的分辨率，调整这个值可以调节速度与精度之间的折衷
model.net_info["height"] = args.reso
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
# 如果有GPU使用，转换模型的数据类型为CUDA
if CUDA:
model.cuda()
from darknet import DarkNet, TinyDarkNet
import cv2
import numpy as np
classes = str('data/dataset-name/coco.names')
classes = load_classes(classes)
num_classes = len(classes)
input_dim = 416
print('number of classes: ',num_classes)
#Get GPU(s) name(es)
ctx = try_gpu([0])[0]
confidence = 0.8
nms_thresh = 0.2
net =DarkNet(input_dim=input_dim,num_classes=num_classes)
anchors = np.array([(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
(59, 119), (116, 90), (156, 198), (373, 326)])
net.initialize(ctx=ctx)
tmp_batch = nd.uniform(shape=(1, 3, input_dim, input_dim), ctx=ctx)
net(tmp_batch)
param=str("./data/weights/yolov3-608.weights")
#Load weights
net.load_weights(param, fine_tune=False)
net.hybridize()
capture = cv2.VideoCapture('./Video/Survei_1.mp4')
while (capture.isOpened()):
shuffle=True)
dataloaders["val"] = val_dataloader
# metrics
obj_loss = LossRecorder('objectness_loss')
cls_loss = LossRecorder('classification_loss')
box_loss = LossRecorder('box_refine_loss')
positive_weight = 1.0
negative_weight = 1.0
# loss with 2x weighting (TODO why?)
l2_loss = L2Loss(weight=2.)
# construct our network, initialize, move to context
net = DarkNet(num_classes=num_classes, input_dim=input_dim)
net.initialize(init=mx.init.Xavier(), ctx=ctx)
# if we provided weights, load them
if args.params is not None:
if args.params.endswith(".params"):
net.load_params(args.params)
elif args.params.endswith(".weights"):
X = nd.uniform(shape=(1, 3, input_dim, input_dim), ctx=ctx[-1])
net(X)
# net.load_weights(args.params, fine_tune=num_classes != 80)
else:
print("params {} load error!".format(args.params))
exit()
print("load params: {}".format(args.params))
nsamples = file_lines(trainlist)
max_epochs = (max_batches*batch_size)//nsamples+1 #why?
print(nsamples)
print(max_epochs)
seed = int(time.time())
torch.manual_seed(seed)
if use_cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = gpus
torch.cuda.manual_seed(seed)
device = torch.device("cuda" if use_cuda else "cpu")
print(device)
model = DarkNet(cfgfile, use_cuda=use_cuda)
if weightfile is not None:
model.load_weights(weightfile)
#model.print_net()
init_epoch = model.seen//nsamples
print(init_epoch)
loss_layers = model.loss_layers
for l in loss_layers:
l.seen = model.seen
if use_cuda:
if ngpus > 1:
model = torch.nn.DataParallel(model).to(device)
batch_size = args.batch_size
confidence = args.confidence
nms_thresh = args.nms_thresh
input_dim = args.input_dim
dst_dir = args.dst_dir
start = 0
classes = load_classes(args.classes)
gpu = [int(x) for x in args.gpu.replace(" ", "").split(",")]
ctx = try_gpu(args.gpu)[0]
num_classes = len(classes)
if args.tiny:
net = TinyDarkNet(input_dim=input_dim, num_classes=num_classes)
anchors = np.array([(10, 14), (23, 27), (37, 58), (81, 82), (135, 169), (344, 319)])
else:
net = DarkNet(input_dim=input_dim, num_classes=num_classes)
anchors = np.array([(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
(59, 119), (116, 90), (156, 198), (373, 326)])
net.initialize(ctx=ctx)
input_dim = args.input_dim
try:
imlist = [os.path.join(images, img) for img in os.listdir(images)]
except NotADirectoryError:
imlist = []
except FileNotFoundError:
print("No file or directory with the name {}".format(images))
if not os.path.exists(dst_dir):
os.mkdir(dst_dir)
elif epoch == 105:
lr = 0.0001
else:
return
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
# Load pre-trained darknet.
darknet = DarkNet(conv_only=True, bn=True, init_weight=True)
darknet.features = torch.nn.DataParallel(darknet.features)
src_state_dict = torch.load(checkpoint_path)['state_dict']
dst_state_dict = darknet.state_dict()
for k in dst_state_dict.keys():
print('Loading weight of', k)
dst_state_dict[k] = src_state_dict[k]
darknet.load_state_dict(dst_state_dict)
# Load YOLO model.
yolo = YOLOv1(darknet.features)
yolo.conv_layers = torch.nn.DataParallel(yolo.conv_layers)
yolo.cuda()
mode="val",
coco_path=args.coco_val)
val_dataloader = gluon.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=True)
dataloaders["val"] = val_dataloader
obj_loss = LossRecorder('objectness_loss')
cls_loss = LossRecorder('classification_loss')
box_loss = LossRecorder('box_refine_loss')
positive_weight = 1.0
negative_weight = 1.0
l2_loss = L2Loss(weight=2.)
net = DarkNet(num_classes=num_classes, input_dim=input_dim)
net.initialize(init=mx.init.Xavier(), ctx=ctx)
if args.params.endswith(".params"):
net.load_params(args.params)
elif args.params.endswith(".weights"):
X = nd.uniform(shape=(1, 3, input_dim, input_dim), ctx=ctx[-1])
net(X)
net.load_weights(args.params, fine_tune=num_classes != 80)
else:
print("params {} load error!".format(args.params))
exit()
print("load params: {}".format(args.params))
net.hybridize()
# for _, w in net.collect_params().items():
# if w.name.find("58") == -1 and w.name.find("66") == -1 and w.name.find("74") == -1:
# w.grad_req = "null"
def main():
# img_path = os.path.join(IMG_DATA, 'hand_and_car.jpg') # 图片路径
img_path = os.path.join(IMG_DATA, 'jiaotong-0727',
'OPjlj2bF9nKjgOg0QNzTxqwRJZdgp.jpg')
params_path = os.path.join(MODEL_DATA, 'yolov3.weights') # YOLO v3 权重文件
classes_path = os.path.join(CONFIGS_DATA, 'coco.names') # 类别文件
confidence = 0.15 # 置信度
nms_thresh = 0.20 # NMS阈值
input_dim = 416 # YOLOv3的检测尺寸
classes_name = load_classes(classes_path) # 加载类别目录
num_classes = len(classes_name) # 类别数
gpu = '1' # GPU
gpu = [int(x) for x in gpu.replace(" ", "").split(",")]
ctx = try_gpu(gpu)[0] # 选择ctx
net = DarkNet(input_dim=input_dim, num_classes=num_classes) # 基础网络DarkNet
anchors = np.array([(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
(59, 119), (116, 90), (156, 198),
(373, 326)]) # anchors
net.initialize(ctx=ctx)
# 加载模型
if params_path.endswith(".params"):
net.load_params(params_path)
elif params_path.endswith(".weights"):
tmp_batch = nd.uniform(shape=(1, 3, input_dim, input_dim), ctx=ctx)
net(tmp_batch)
net.load_weights(params_path, fine_tune=False)
else:
print("params {} load error!".format(params_path))
exit()
print("load params: {}".format(params_path))
net.hybridize()
image_data = cv2.imread(img_path) # 读取图片数据
image_reform = reform_img(image_data, input_dim)
image_arr = nd.array([image_reform], ctx=ctx)
prediction = predict_transform(net(image_arr), input_dim, anchors)
pred_res = filter_results(prediction,
num_classes,
confidence=confidence,
nms_conf=nms_thresh)
boxes, scores, classes = generate_bboxes(image_data,
pred_res,
input_dim=input_dim)
hsv_tuples = [(float(x) / len(classes_name), 1., 1.)
for x in range(len(classes_name))] # 不同颜色
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
colors)) # RGB
np.random.seed(10101)
np.random.shuffle(colors) # 随机颜色
np.random.seed(None)
img_data = Image.open(img_path)
image = draw_boxes(img_data, boxes, scores, classes, colors, classes_name)
image.show()
print(pred_res)
def main_worker(gpu, ngpus_per_node, writer, log_dir, args):
global best_acc1
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
#if args.pretrained:
# print("=> using pre-trained model '{}'".format(args.arch))
# model = models.__dict__[args.arch](pretrained=True)
#else:
# print("=> creating model '{}'".format(args.arch))
# model = models.__dict__[args.arch]()
# create model
print("=> creating model DarkNet (bn: {}) ...".format(args.bn))
model = DarkNet(bn=args.bn)
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(args.workers / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
#if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
# model.features = torch.nn.DataParallel(model.features)
# model.cuda()
#else:
# model = torch.nn.DataParallel(model).cuda()
model.features = torch.nn.DataParallel(model.features)
model.cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, None, args)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, writer, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, writer, args)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, epoch, writer, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
#'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
},
is_best,
log_dir)
mode="val",
coco_path=None)
val_dataloader = gluon.data.DataLoader(val_dataset,
batch_size=batch_size,
shuffle=True)
dataloaders["val"] = val_dataloader
obj_loss = LossRecorder('objectness_loss')
cls_loss = LossRecorder('classification_loss')
box_loss = LossRecorder('box_refine_loss')
positive_weight = 1.0
negative_weight = 1.0
l2_loss = L2Loss(weight=2.)
net = DarkNet(num_classes=num_classes, input_dim=input_dim)
net.initialize(init=mx.init.Xavier(), ctx=ctx) # 初始化神经网络
if args.params.endswith(".params"):
net.load_params(args.params)
elif args.params.endswith(".weights"):
X = nd.uniform(shape=(1, 3, input_dim, input_dim),
ctx=ctx) # feed shape tip here
net(X) # debuging what is it ?
#print("params {} loading ......".format(args.params))
#net.load_weights(args.params, fine_tune=num_classes != 80)
else:
print("params {} load error!".format(args.params))
exit()
net.hybridize()
try:
classes = load_classes('../../data/coco.names')
except FileNotFoundError:
classes = load_classes('data/coco.names')
no_classes = len(classes)
params = cmd_line()
imgs = params.imgs
nms = float(params.nms)
batch_size = int(params.bs)
conf = float(params.conf)
gpu = torch.cuda.is_available()
model = DarkNet(params.cfg) #setup
model.load_weights(params.weights)
print("network loaded")
if gpu:
model.cuda()
model.network_info['height'] = params.reso
input_dim = int(model.network_info['height'])
if ((input_dim % 32 != 0) or (input_dim <= 32)):
raise Exception('invalid image size')
model.eval() #eval mode
read_s = time.time() #need to note time
try: #get image locations in list
image_paths = [
imgs = parse.imgs
outdet = parse.det
cfg = parse.cfg
weights = parse.weights
batch_size = int(parse.bs)
confidence = float(parse.conf)
nms_threshold = float(parse.nms_thre)
resolution = parse.res
CUDA = torch.cuda.is_available()
classes = load_names("./data/coco.names")
num_class = 80
# set up neural networks
print("Loading Network...")
model = DarkNet(cfg)
model.load_weights(weights)
print("Finish loading")
model.net_info["height"] = resolution
img_dim = int(resolution)
assert img_dim >= 32
assert img_dim % 32 == 0
if CUDA:
model.cuda()
model.eval()
t_read_dir = time.time()
# get image list, 所有需要测试的图片集
