def __init__(self, config, is_training=True):
super(ModelMain, self).__init__()
self.config = config
self.training = is_training
self.model_params = config["model_params"]
# backbone
# _backbone_fn = backbone_fn[self.model_params["backbone_name"]]
# self.backbone = _backbone_fn(self.model_params["backbone_pretrained"])
# self.backbone = mobilenet.mobilenetv2(self.model_params["backbone_pretrained"])
self.backbone = ShuffleNetV2(scale=1,
in_channels=3,
c_tag=0.5,
num_classes=2,
activation=nn.ReLU,
SE=False,
residual=False)
_out_filters = self.backbone.layers_out_filters
# embedding0
final_out_filter0 = len(
config["yolo"]["anchors"][0]) * (5 + config["yolo"]["classes"])
self.embedding0 = self._make_embedding([512, 1024], _out_filters[-1],
final_out_filter0)
# embedding1
final_out_filter1 = len(
config["yolo"]["anchors"][1]) * (5 + config["yolo"]["classes"])
self.embedding1_cbl = self._make_cbl(512, 256, 1)
self.embedding1_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.embedding1 = self._make_embedding([256, 512],
_out_filters[-2] + 256,
final_out_filter1)
# embedding2
final_out_filter2 = len(
config["yolo"]["anchors"][2]) * (5 + config["yolo"]["classes"])
self.embedding2_cbl = self._make_cbl(256, 128, 1)
self.embedding2_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.embedding2 = self._make_embedding([128, 256],
_out_filters[-3] + 128,
final_out_filter2)
self.yolo_losses = []
self.losses_name = [
"total_loss", "x", "y", "w", "h", "conf", "cls", "recall"
]
if not self.training:
self.nms = NMSLayer()
for i in range(3):
yololayer = YOLOLayer(config["batch_size"], i,
config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"]))
self.yolo_losses.append(yololayer)
def build_yolov3(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
net = nn.DataParallel(net)
net = net.cuda()
yolo_losses = [] # YOLO loss with 3 scales
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
return net, yolo_losses
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
ini_files = [
inifile for inifile in os.listdir(
os.path.join(config['test_weights'], 'result'))
if inifile.endswith('.ini')
]
accuracy_s = [(inifile[:-4]).split('_')[-1] for inifile in ini_files]
accuracy_ints = list(map(float, accuracy_s))
max_index = accuracy_ints.index(max(accuracy_ints))
# for kkk,ini_file in enumerate(ini_files):
ini_list_config = configparser.ConfigParser()
config_file_path = os.path.join(config['test_weights'], 'result',
ini_files[max_index])
Bi_picpath = os.path.join(config['test_weights'], 'result',
ini_files[max_index]).replace('.ini', '')
os.makedirs(Bi_picpath, exist_ok=True)
ini_list_config.read(config_file_path)
ini_session = ini_list_config.sections()
accuracy = ini_list_config.items(ini_session[0])
err_jpgfiles = ini_list_config.items(ini_session[1])
weight_file = os.path.join(
config['test_weights'],
'%s.weights' % ini_files[max_index].split('_')[0])
if weight_file: # Restore pretrain model
logging.info("load checkpoint from {}".format(weight_file))
state_dict = torch.load(weight_file)
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# images_name = os.listdir(config["images_path"]) # prepare images path
# images_path = [os.path.join(config["images_path"], name) for name in images_name]
# if len(images_path) == 0:
# raise Exception("no image found in {}".format(config["images_path"]))
# batch_size = config["batch_size"]# Start inference
# for step in range(0, len(images_path), batch_size):
for _jpg_images in err_jpgfiles:
images = [] # preprocess
images_origin = []
jpg_path = str(_jpg_images[1])
logging.info("processing: {}".format(jpg_path))
bbox_list = read_gt_boxes(jpg_path)
image = cv2.imread(jpg_path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(jpg_path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
with torch.no_grad(): # inference
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
classes = open(config["classes_names_path"],
"r").read().split("\n")[:-1]
for idx, detections in enumerate(batch_detections):
image_show = images_origin[idx]
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
ori_h, ori_w = images_origin[
idx].shape[:
2] # Rescale coordinates to original dimensions
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
image_show = cv2.rectangle(images_origin[idx], (x1, y1),
(x1 + box_w, y1 + box_h),
(0, 255, 0), 2)
for (x1, x2, y1, y2) in bbox_list:
[x1, x2, y1, y2] = map(int, [x1, x2, y1, y2])
cv2.rectangle(image_show, (x1, y1), (x2, y2), (255, 0, 0),
2)
pic_name = (jpg_path.split('/')[-1]).split('.')[0]
image_show = cv2.cvtColor(image_show, cv2.COLOR_RGB2BGR)
cv2.imwrite(
os.path.join(Bi_picpath,
'%s.jpg' % os.path.basename(pic_name)),
image_show)
def train(config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain_SN(config, is_training=is_training)
net.train(is_training)
# Optimizer and learning rate
optimizer = _get_optimizer(config, net)
lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=config["lr"]["decay_step"],
gamma=config["lr"]["decay_gamma"])
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("Load pretrained weights from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(
config["train_path"], (config["img_w"], config["img_h"]),
is_training=True,
is_scene=False),
batch_size=config["batch_size"],
shuffle=True,
drop_last=True,
num_workers=0,
pin_memory=True)
# Start the training loop
logging.info("Start training.")
dataload_len = len(dataloader)
for epoch in range(config["epochs"]):
recall = 0
mini_step = 0
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
start_time = time.time()
config["global_step"] += 1
for mini_batch in range(8):
mini_step += 1
# Forward and backward
optimizer.zero_grad()
outputs = net(images)
losses_name = [
"total_loss", "x", "y", "w", "h", "conf", "cls", "recall"
]
losses = [0] * len(losses_name)
for i in range(3):
_loss_item = yolo_losses[i](outputs[i], labels)
for j, l in enumerate(_loss_item):
losses[j] += l
# losses = [sum(l) for l in losses]
loss = losses[0]
loss.backward()
optimizer.step()
_loss = loss.item()
# example_per_second = config["batch_size"] / duration
# lr = optimizer.param_groups[0]['lr']
strftime = datetime.datetime.now().strftime("%H:%M:%S")
if (losses[7] / 3 >= recall / (step + 1)) or mini_batch == 7:
recall += losses[7] / 3
print(
'%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %d]'
%
(strftime, epoch, config["epochs"], step, dataload_len,
losses[1], losses[2], losses[3], losses[4], losses[5],
losses[6], _loss, losses[7] / 3, recall /
(step + 1), mini_batch))
break
else:
print(
'%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,prerc %.3f %d]'
% (strftime, epoch, config["epochs"], step,
dataload_len, losses[1], losses[2], losses[3],
losses[4], losses[5], losses[6], _loss,
losses[7] / 3, recall / step, mini_batch))
if (epoch % 2 == 0 and recall / len(dataloader) > 0.7
) or recall / len(dataloader) > 0.96:
torch.save(
net.state_dict(), '%s/%.4f_%04d.weights' %
(checkpoint_dir, recall / len(dataloader), epoch))
lr_scheduler.step()
# net.train(True)
logging.info("Bye bye")
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# prepare images path
images_path = os.listdir(config["images_path"])
images_path = [file for file in images_path if file.endswith('.jpg')]
# images_path = [os.path.join(config["images_path"], name) for name in images_name]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
bgimage = cv2.imread(os.path.join(config["images_path"], images_path[0]), cv2.IMREAD_COLOR)
bgimage = cv2.cvtColor(bgimage, cv2.COLOR_BGR2GRAY)
for step in range(0, len(images_path)-1, batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step*batch_size: (step+1)*batch_size]:
if not path.endswith(".jpg") and (not path.endswith(".png")) and not path.endswith(".JPEG"):
continue
image = cv2.imread(os.path.join(config["images_path"], path), cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
for idx, detections in enumerate(batch_detections):
image_show =images_origin[idx]
if detections is not None:
anno = savexml.GEN_Annotations(path + '.jpg')
anno.set_size(1280, 720, 3)
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_list = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox_list.append((x1, y1,box_w,box_h))
image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 0, 255), 1)
boundbox = bg_judge(images_origin[idx],bgimage,bbox_list)
print('boundbox',boundbox,bbox_list)
for (x,y,w,h) in boundbox:
image_show = cv2.rectangle(image_show, (x, y), (x + w, y + h), (0, 255, 0), 1)
# anno.add_pic_attr("mouse", int(x1.cpu().data), int(y1.cpu().data), int(box_w.cpu().data) , int(box_h.cpu().data) ,"0")
#
# xml_path = os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\Annotations').replace('jpg','xml')
# anno.savefile(xml_path)
# cv2.imwrite(os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\rec_pic'),images_origin[idx])
cv2.imshow('1', image_show)
cv2.waitKey(1)
logging.info("Save all results to ./output/")
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"],
nms_thres=config["iou_thres"])
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"],
"r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(images_origin[idx])
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_colors = random.sample(colors, n_cls_preds)
boxes = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
color = bbox_colors[int(
np.where(unique_labels == int(cls_pred))[0])]
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
box = BoundBox(x1, y1, x1 + box_w, y1 + box_h,
cls_conf.item(), int(cls_pred),
classes[int(cls_pred)])
boxes.append(box)
# bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=1,
# edgecolor=color,
# facecolor='none')
# # Add the bbox to the plot
# ax.add_patch(bbox)
# Add label
# plt.text(x1, y1, s=classes[int(cls_pred)], color='white',
# verticalalignment='top',
# bbox={'color': color, 'pad': 0})
# Save generated image with detections
img_show = draw_boxes(images_origin[idx], boxes, labels, 0.5)
img_show = cv2.resize(img_show,
(img_show.shape[1], img_show.shape[0]),
interpolation=cv2.INTER_CUBIC)
# outVideo.write(img_show)
cv2.imshow("ai", img_show)
cv2.waitKey()
# plt.close()
logging.info("Save all results to ./output/")
index += 1
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
dataloader = torch.utils.data.DataLoader(COCODataset(
config["train_path"], (config["img_w"], config["img_h"]),
is_training=True,
is_scene=True),
batch_size=config["batch_size"],
shuffle=True,
drop_last=True,
num_workers=0,
pin_memory=True)
def validate(net):
n_gt = 0
def train(config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Optimizer and learning rate
optimizer = _get_optimizer(config, net)
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=15)
# lr_scheduler = optim.lr_scheduler.StepLR(
# optimizer,
# step_size=config["lr"]["decay_step"],
# gamma=config["lr"]["decay_gamma"])
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("Load pretrained weights from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
# Only export onnx
# if config.get("export_onnx"):
# real_model = net.module
# real_model.eval()
# dummy_input = torch.randn(8, 3, config["img_h"], config["img_w"]).cuda()
# save_path = os.path.join(config["sub_working_dir"], "pytorch.onnx")
# logging.info("Exporting onnx to {}".format(save_path))
# torch.onnx.export(real_model, dummy_input, save_path, verbose=False)
# logging.info("Done. Exiting now.")
# sys.exit()
# Evaluate interface
# if config["evaluate_type"]:
# logging.info("Using {} to evaluate model.".format(config["evaluate_type"]))
# evaluate_func = importlib.import_module(config["evaluate_type"]).run_eval
# config["online_net"] = net
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(
config["train_path"], (config["img_w"], config["img_h"]),
is_training=True,
is_scene=True),
batch_size=config["batch_size"],
shuffle=True,
drop_last=True,
num_workers=0,
pin_memory=True)
# Start the training loop
logging.info("Start training.")
dataload_len = len(dataloader)
best_acc = 0.5
for epoch in range(config["epochs"]):
recall = 0
mini_step = 0
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
config["global_step"] += 1
# Forward and backward
optimizer.zero_grad()
outputs = net(images.cuda())
losses_name = [
"total_loss", "x", "y", "w", "h", "conf", "cls", "recall"
]
losses = [0] * len(losses_name)
for i in range(3):
_loss_item = yolo_losses[i](outputs[i], labels)
for j, l in enumerate(_loss_item):
losses[j] += l
# losses = [sum(l) for l in losses]
loss = losses[0]
loss.backward()
optimizer.step()
_loss = loss.item()
# example_per_second = config["batch_size"] / duration
lr = optimizer.param_groups[0]['lr']
strftime = datetime.datetime.now().strftime("%H:%M:%S")
# if (losses[7] / 3 >= recall / (step + 1)):#mini_batch为0走这里
recall += losses[7] / 3
print(
'%s [Epoch %d/%d,batch %03d/%d loss:x %.5f,y %.5f,w %.5f,h %.5f,conf %.5f,cls %.5f,total %.5f,rec %.3f,avrec %.3f %.3f]'
% (strftime, epoch, config["epochs"], step, dataload_len,
losses[1], losses[2], losses[3], losses[4], losses[5],
losses[6], _loss, losses[7] / 3, recall / (step + 1), lr))
if recall / len(dataloader) > best_acc:
best_acc = recall / len(dataloader)
if epoch > 0:
torch.save(
net.state_dict(), '%s/%.4f_%04d.weights' %
(checkpoint_dir, recall / len(dataloader), epoch))
lr_scheduler.step()
net.train(is_training)
torch.cuda.empty_cache()
# net.train(True)
logging.info("Bye bye")
def evaluate(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(
config["val_path"], (config["img_w"], config["img_h"]), is_training=1),
batch_size=config["batch_size"],
shuffle=False,
drop_last=True,
num_workers=0,
pin_memory=False)
# Start the eval loop
logging.info("Start eval.")
n_gt = 0
correct = 0
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
labels = labels.cuda()
with torch.no_grad():
time1 = datetime.datetime.now()
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
output = non_max_suppression(output, 1, conf_thres=0.3)
print("time", (datetime.datetime.now() - time1).microseconds)
# calculate
for sample_i in range(labels.size(0)):
# Get labels for sample where width is not zero (dummies)
target_sample = labels[sample_i, labels[sample_i, :, 3] != 0]
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
tx1, tx2 = config["img_w"] * (
tx - tw / 2), config["img_w"] * (tx + tw / 2)
ty1, ty2 = config["img_h"] * (
ty - th / 2), config["img_h"] * (ty + th / 2)
n_gt += 1
box_gt = torch.cat([
coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]
]).view(1, -1)
sample_pred = output[sample_i]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[
sample_pred[:, 6] == obj_cls]:
box_pred = torch.cat([
coord.unsqueeze(0)
for coord in [x1, y1, x2, y2]
]).view(1, -1)
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thres"]:
correct += 1
break
if n_gt:
logging.info('Batch [%d/%d] mAP: %.5f' %
(step, len(dataloader), float(correct / n_gt)))
logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def test(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start testing FPS of different batch size
for batch_size in range(1, 10):
# preprocess
images = []
for path in images_path[:batch_size]:
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
for i in range(batch_size - len(images)):
images.append(images[0]) # fill len to batch_sze
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference in 30 times and calculate average
inference_times = []
for i in range(30):
start_time = time.time()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
torch.cuda.synchronize() # wait all done.
end_time = time.time()
inference_times.append(end_time - start_time)
inference_time = sum(inference_times) / len(
inference_times) / batch_size
fps = 1.0 / inference_time
logging.info(
"Batch_Size: {}, Inference_Time: {:.5f} s/image, FPS: {}".format(
batch_size, inference_time, fps))
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [os.path.join(config["images_path"], name) for name in images_name]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture("./007.avi")
img_i = 0
start = time.time()
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
img_i += 1
# preprocess
images = []
images_origin = []
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
time1=datetime.datetime.now()
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
print("time1",(datetime.datetime.now()-time1).microseconds)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
print("time2", (datetime.datetime.now() - time1).microseconds)
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
img_show = images_origin[idx]
img_show = cv2.cvtColor(img_show, cv2.COLOR_RGB2BGR)
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
boxes=[]
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
box = BoundBox(x1, y1, x1 + box_w, y1 + box_h, cls_conf.item(), int(cls_pred))
boxes.append(box)
img_show = draw_boxes(img_show, boxes, labels)
# image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0), 1)
cv2.imshow('1', img_show)
cv2.waitKey(1)
logging.info("Save all results to ./output/")
def train(config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Optimizer and learning rate
optimizer = _get_optimizer(config, net)
lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=config["lr"]["decay_step"],
gamma=config["lr"]["decay_gamma"])
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("Load pretrained weights from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
total_loss = 0
last_total_loss = 0
manager = Manager()
# 父进程创建Queue，并传给各个子进程：
q = manager.Queue(1)
lock = manager.Lock() # 初始化一把锁
p = Pool()
pw = p.apply_async(get_data, args=(q, lock))
batch_len = q.get()
if batch_len[0] == "len":
batch_len = batch_len[1]
logging.info("Start training.")
for epoch in range(config["epochs"]):
recall = 0
for step in range(batch_len):
samples = q.get()
images, labels = samples["image"], samples["label"]
start_time = time.time()
config["global_step"] += 1
# Forward and backward
optimizer.zero_grad()
outputs = net(images)
losses_name = [
"total_loss", "x", "y", "w", "h", "conf", "cls", "recall"
]
losses = [0] * len(losses_name)
for i in range(3):
_loss_item = yolo_losses[i](outputs[i], labels)
for j, l in enumerate(_loss_item):
losses[j] += l
# losses = [sum(l) for l in losses]
loss = losses[0]
loss.backward()
optimizer.step()
if step > 0 and step % 2 == 0:
_loss = loss.item()
duration = float(time.time() - start_time)
example_per_second = config["batch_size"] / duration
lr = optimizer.param_groups[0]['lr']
strftime = datetime.datetime.now().strftime("%H:%M:%S")
recall += losses[7] / 3
print(
'%s [Epoch %d/%d, Batch %03d/%d losses: x %.5f, y %.5f, w %.5f, h %.5f, conf %.5f, cls %.5f, total %.5f, recall: %.3f]'
% (strftime, epoch, config["epochs"], step, batch_len,
losses[1], losses[2], losses[3], losses[4], losses[5],
losses[6], _loss, losses[7] / 3))
# logging.info(epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f "%
# (epoch, step, _loss, example_per_second, lr))
# config["tensorboard_writer"].add_scalar("lr",
# lr,
# config["global_step"])
# config["tensorboard_writer"].add_scalar("example/sec",
# example_per_second,
# config["global_step"])
# for i, name in enumerate(losses_name):
# value = _loss if i == 0 else losses[i]
# config["tensorboard_writer"].add_scalar(name,
# value,
# config["global_step"])
if (epoch % 2 == 0
and recall / batch_len > 0.7) or recall / batch_len > 0.96:
torch.save(net.state_dict(),
'%s/%04d.weights' % (checkpoint_dir, epoch))
lr_scheduler.step()
def test(config,int_dir='result'):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
ini_files = os.listdir(os.path.join(config['test_weights'], int_dir))
for kkk,ini_file in enumerate(ini_files):
ini_list_config = configparser.ConfigParser()
config_file_path = os.path.join(config['test_weights'], int_dir,ini_files[-kkk-1])
ini_list_config.read(config_file_path)
ini_session = ini_list_config.sections()
# accuracy = ini_list_config.items(ini_session[0])
err_jpgfiles = ini_list_config.items(ini_session[1])
aaa = glob.glob(os.path.join(config['test_weights'],'*_%s.weights'%ini_files[-kkk-1].split('_')[-1].split('.')[0]))
weight_file = aaa[0]#os.path.join(config['test_weights'],'%s.weights'%ini_files[-kkk-1].split('_')[0])
if weight_file: # Restore pretrain model
logging.info("load checkpoint from {}".format(weight_file))
state_dict = torch.load(weight_file)
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(1, i, config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
for index, _jpg_images in enumerate(err_jpgfiles):
images = []# preprocess
images_origin = []
jpg_path = str(_jpg_images[1])
print(str(index+1),jpg_path)
bbox_list = read_gt_boxes(jpg_path)
image = cv2.imread(jpg_path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(jpg_path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
with torch.no_grad():# inference
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
image_show=images_origin[idx]
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
ori_h, ori_w = images_origin[idx].shape[:2]# Rescale coordinates to original dimensions
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
#绿色代表预测，红色代表标注
image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0),2)
for (x1, x2, y1, y2) in bbox_list:
[x1, x2, y1, y2] = map(int, [x1, x2, y1, y2])
cv2.rectangle(image_show, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.imshow('1', image_show)
cv2.waitKey()