def detect(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"]:
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(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# Load tested img
imgfile = config["img_path"]
img = Image.open(imgfile).convert('RGB')
resized = img.resize((config["img_w"], config["img_h"]))
input = image2torch(resized)
input = input.to(torch.device("cuda"))
start = time.time()
outputs = net(input)
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,
config["yolo"]["classes"],
conf_thres=0.5,
nms_thres=0.4)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
namefile = config["classname_path"]
class_names = load_class_names(namefile)
plot_boxes(img, output, 'predictions.jpg', class_names)
def initial_yolo_model(config, size):
is_training = False
config["img_w"] = size
config["img_h"] = size
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
if torch.cuda.is_available():
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
if gpu:
state_dict = torch.load(config["pretrain_snapshot"])
else:
state_dict = torch.load(config["pretrain_snapshot"], map_location=torch.device('cpu'))
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
return net
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 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()
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(
YOLOLoss(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"])
# 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)
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
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
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
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}_{}.jpg'.format(step, idx),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
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.CosineAnnealingLR(optimizer, T_max=10)
# 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=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"]
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()
_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 main(video_fn):
logging.basicConfig(level=logging.DEBUG,
format="[%(asctime)s %(filename)s] %(message)s")
if len(sys.argv) != 2:
logging.error("Usage: python video.py params.py")
sys.exit()
params_path = sys.argv[1]
if not os.path.isfile(params_path):
logging.error("no params file found! path: {}".format(params_path))
sys.exit()
config = importlib.import_module(params_path[:-3]).TRAINING_PARAMS
config["batch_size"] *= len(config["parallels"])
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()
# load pretrained 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(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# load class names
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
cap = cv2.VideoCapture(video_fn)
# Check if camera opened successfully
if (cap.isOpened() == False):
print("Error opening video stream or file")
# Read until video is completed
while (cap.isOpened()):
# Capture frame-by-frame
ret, frame = cap.read()
frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
if ret == True:
# 1. pre-process image
logging.info("processing frame")
image_tensor = prep_image(frame, config)
with torch.no_grad():
outputs = net(image_tensor)
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=0.45)
for idx, detections in enumerate(batch_detections):
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)
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
x1, y1, box_w, box_h = get_rescaled_coords(
frame.shape[0], frame.shape[1], config["img_h"],
config["img_w"], x1, y1, x2, y2)
cv2.rectangle(frame, (x1, y1),
(x1 + box_w, y1 + box_h), color, 2)
cv2.putText(frame, classes[int(cls_pred)], (x1, y1),
cv2.FONT_HERSHEY_SIMPLEX, 1, color, 1,
cv2.LINE_AA)
cv2.imshow('Frame', frame)
# Press Q on keyboard to exit
if cv2.waitKey(25) & 0xFF == ord('q'):
break
# Break the loop
else:
break
# When everything done, release the video capture object
cap.release()
# Closes all the frames
cv2.destroyAllWindows()
def train(config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
# 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)
# Start the training loop
logging.info("Start training.")
dataload_len = len(dataloader)
epoch_size = 4
start = time.time()
pruned_pct = 0
global index, pruned_book, num_pruned
global num_weights
global weight_masks, bias_masks
for epoch in range(config["epochs"]):
if epoch % 4 == 0:
index = 0
num_pruned = 0
num_weights = 0
net.apply(prune)
torch.save(net.state_dict(),
'%s/%.4f_%04d.weights' % (checkpoint_dir, 0.01, 1))
print('previously pruned: %.3f %%' % (100 * (pruned_pct)))
print('number pruned: %.3f %%' % (100 *
(num_pruned / num_weights)))
new_pruned = num_pruned / num_weights - pruned_pct
pruned_pct = num_pruned / num_weights
# if new_pruned <= 0.01:
# time_elapse = time.time() - start
# print('training time:', str(timedelta(seconds=time_elapse)))
# break
recall = 0
mini_step = 0
for step, samples in enumerate(dataloader):
index = 0
images, labels = samples["image"], samples["label"]
start_time = time.time()
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()
net.apply(set_grad)
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")
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 (epoch % 2 == 0 and recall / len(dataloader) > 0.5
) or recall / len(dataloader) > 0:
# torch.save(net.state_dict(), '%s/%.4f_%04d.weights' % (checkpoint_dir, recall / len(dataloader), epoch))
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 evaluate(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: {}".format(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(
YOLOLoss(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=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=False)
# Coco Prepare.
index2category = json.load(open("coco_index2category.json"))
# Start the eval loop
logging.info("Start eval.")
coco_results = []
coco_img_ids = set([])
APs = []
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
image_paths, origin_sizes = samples["image_path"], samples[
"origin_size"]
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=0.0001,
nms_thres=0.45)
for idx, detections in enumerate(batch_detections):
correct = []
annotations = labels[idx, labels[idx, :, 3] != 0]
image_id = int(os.path.basename(image_paths[idx])[-16:-4])
coco_img_ids.add(image_id)
if detections is None:
if annotations.size(0) != 0:
APs.append(0)
continue
detections = detections[np.argsort(-detections[:, 4])]
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
# ===========================================================================================================================
# The amount of padding that was added
pad_x = max(origin_size[1] - origin_size[0],
0) * (config["img_w"] / max(origin_size))
pad_y = max(origin_size[0] - origin_size[1],
0) * (config["img_w"] / max(origin_size))
# Image height and width after padding is removed
unpad_h = config["img_w"] - pad_y
unpad_w = config["img_w"] - pad_x
# ===========================================================================================================================
if annotations.size(0) == 0:
correct.extend([0 for _ in range(len(detections))])
else:
target_boxes = torch.FloatTensor(annotations[:, 1:].shape)
target_boxes[:,
0] = (annotations[:, 1] - annotations[:, 3] / 2)
target_boxes[:,
1] = (annotations[:, 2] - annotations[:, 4] / 2)
target_boxes[:,
2] = (annotations[:, 1] + annotations[:, 3] / 2)
target_boxes[:,
3] = (annotations[:, 2] + annotations[:, 4] / 2)
target_boxes *= config["img_w"]
detected = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
pred_bbox = (x1, y1, x2, y2)
#x1 = x1 / config["img_w"] * origin_size[0]
#x2 = x2 / config["img_w"] * origin_size[0]
#y1 = y1 / config["img_h"] * origin_size[1]
#y2 = y2 / config["img_h"] * origin_size[1]
#w = x2 - x1
#h = y2 - y1
h = ((y2 - y1) / unpad_h) * origin_size[1]
w = ((x2 - x1) / unpad_w) * origin_size[0]
y1 = ((y1 - pad_y // 2) / unpad_h) * origin_size[1]
x1 = ((x1 - pad_x // 2) / unpad_w) * origin_size[0]
coco_results.append({
"image_id":
image_id,
"category_id":
index2category[str(int(cls_pred.item()))],
"bbox": (float(x1), float(y1), float(w), float(h)),
"score":
float(conf),
})
pred_bbox = torch.FloatTensor(pred_bbox).view(1, -1)
# Compute iou with target boxes
iou = bbox_iou(pred_bbox, target_boxes)
# Extract index of largest overlap
best_i = np.argmax(iou)
# If overlap exceeds threshold and classification is correct mark as correct
if iou[best_i] > config[
'iou_thres'] and cls_pred == annotations[
best_i, 0] and best_i not in detected:
correct.append(1)
detected.append(best_i)
else:
correct.append(0)
true_positives = np.array(correct)
false_positives = 1 - true_positives
# Compute cumulative false positives and true positives
false_positives = np.cumsum(false_positives)
true_positives = np.cumsum(true_positives)
# Compute recall and precision at all ranks
recall = true_positives / annotations.size(0) if annotations.size(
0) else true_positives
precision = true_positives / np.maximum(
true_positives + false_positives,
np.finfo(np.float64).eps)
# Compute average precision
AP = compute_ap(recall, precision)
APs.append(AP)
print("+ Sample [%d/%d] AP: %.4f (%.4f)" %
(len(APs), 5000, AP, np.mean(APs)))
logging.info("Now {}/{}".format(step, len(dataloader)))
print("Mean Average Precision: %.4f" % np.mean(APs))
save_results_path = "coco_results.json"
with open(save_results_path, "w") as f:
json.dump(coco_results,
f,
sort_keys=True,
indent=4,
separators=(',', ':'))
logging.info("Save coco format results to {}".format(save_results_path))
# COCO api
logging.info("Using coco-evaluate tools to evaluate.")
cocoGt = COCO(config["annotation_path"])
cocoDt = cocoGt.loadRes(save_results_path)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.params.imgIds = list(coco_img_ids) # real imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
def train(imgs, labels, checkpoint_path, config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
# 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 checkpoint_path:
logging.info("Load pretrained weights from {}".format(checkpoint_path))
state_dict = torch.load(checkpoint_path)
net.load_state_dict(state_dict)
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(SatDataset(
imgs, labels, (config["img_w"], config["img_h"]), is_training=True),
batch_size=config["batch_size"],
shuffle=True,
num_workers=1,
pin_memory=True)
# Start the training loop
logging.info("Start training.")
for epoch in range(config["epochs"]):
for step, samples in enumerate(dataloader):
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"]
losses = [[]] * 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].append(l)
losses = [sum(l) for l in losses]
loss = losses[0]
loss.backward()
optimizer.step()
if step > 0 and step % 10 == 0:
_loss = loss.item()
duration = float(time.time() - start_time)
example_per_second = config["batch_size"] / duration
lr = optimizer.param_groups[0]['lr']
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"])
lr_scheduler.step()
# net.train(False)
checkpoint_path = _save_checkpoint(net.state_dict(), config)
# net.train(True)
logging.info("Bye~")
return checkpoint_path
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"])
# 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):
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),
classes[int(cls_pred)])
boxes.append(box)
# 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()
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)
t_max = 50
# lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t_max,eta_min=1e-05)
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
# 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"] * config["parallels"],
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.2
last_recall = 0.6
for epoch in range(config["epochs"]):
recall = 0
mini_step = 0
for step, samples in enumerate(dataloader):
start = time.time()
images, labels = samples["image"], samples["label"]
config["global_step"] += 1
# Forward and backward
optimizer.zero_grad()
losses = net(images.cuda(), labels.cuda())
# current_recall = mAP(detections, labels, config["img_w"])
# current_recall = np.mean(current_recall)
if config["parallels"] > 1:
losses = losses.view(config["parallels"], 8)[0] + losses.view(
config["parallels"], 8)[1]
loss = losses[0]
if epoch > 0:
loss = loss * 20
current_recall = float(losses[7] / 3 / config["parallels"])
if last_recall < 0.65:
loss = loss + 20 * (1 - current_recall) # * 0.8
else:
loss = loss + 20 * (1 - current_recall)
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 += current_recall
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, current_recall, recall / (step + 1), lr))
last_recall = recall / len(dataloader)
if recall / len(dataloader) > best_acc:
best_acc = recall / len(dataloader)
torch.save(
net.state_dict(), '%s/%.4f_%04d.weights' %
(checkpoint_dir, recall / len(dataloader), epoch))
lr_scheduler.step()
# if epoch % (lr_scheduler.T_max + next_need) == (lr_scheduler.T_max + next_need - 1):
# next_need += float(lr_scheduler.T_max)
# lr_scheduler.T_max += 2
# lr_scheduler.last_epoch = 0
# lr_scheduler.base_lrs*=0.98
# lr_scheduler.base_lrs[0] *= 0.95
# lr_scheduler.base_lrs[1] *= 0.95
# net.train(is_training)
# torch.cuda.empty_cache()
# net.train(True)
logging.info("Bye bye")
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):
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 evaluate(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"]:
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(YOLOLoss(config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(dataset=COCODataset(config["val_path"], config["img_w"]),
batch_size=config["batch_size"],
shuffle=True, num_workers=1, pin_memory=False)
# Start the eval loop
logging.info("Start eval.")
n_gt = 0
correct = 0
logging.info('%s' % str(dataloader))
gt_histro={}
pred_histro = {}
correct_histro = {}
for i in range(config["yolo"]["classes"]):
gt_histro[i] = 1
pred_histro[i] = 1
correct_histro[i] = 0
# images 是一个batch里的全部图片，labels是一个batch里面的全部标签
for step, (images, labels) in enumerate(dataloader):
labels = labels.cuda()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
# 把三个尺度上的预测结果在第1维度(第0维度是batch里的照片，第1维度是一张照片里面的各个预测框，第2维度是各个预测数值)上拼接起来
batch_output = torch.cat(output_list, dim=1)
logging.info('%s' % str(batch_output.shape))
# 进行NMS抑制
batch_output = non_max_suppression(prediction=batch_output, num_classes=config["yolo"]["classes"], conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"])
# calculate
for sample_index_in_batch in range(labels.size(0)):
# fetched img sample in tensor( C(RxGxB) x H x W ), transform to cv2 format in H x W x C(BxGxR)
sample_image = images[sample_index_in_batch].numpy()
sample_image = np.transpose(sample_image, (1, 2, 0))
sample_image = cv2.cvtColor(sample_image, cv2.COLOR_RGB2BGR)
logging.debug("fetched img %d size %s" % (sample_index_in_batch, sample_image.shape))
# Get labels for sample where width is not zero (dummies)(init all labels to zeros in array)
target_sample = labels[sample_index_in_batch, labels[sample_index_in_batch, :, 3] != 0]
# get prediction for this sample
sample_pred = batch_output[sample_index_in_batch]
if sample_pred is not None:
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred: # for each prediction box
# logging.info("%d" % obj_cls)
box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1)
sample_image = draw_prediction(sample_image,conf, obj_conf, int(obj_pred), (x1, y1, x2, y2), config)
# 每一个ground truth的 分类编号obj_cls、相对中心x、相对中心y、相对宽w、相对高h
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
# 转化为输入像素尺寸的 左上角像素tx1 ty1，右下角像素tx2 ty2
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)
# 计算ground truth数量，用于统计信息
n_gt += 1
gt_histro[int(obj_cls)] += 1
# 转化为 shape(1,4)的tensor，用来计算IoU
box_gt = torch.cat([coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]]).view(1, -1)
# logging.info('%s' % str(box_gt.shape))
sample_pred = batch_output[sample_index_in_batch]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
# 对于每一个ground truth，遍历预测结果
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[sample_pred[:, 6] == obj_cls]: # 如果当前预测分类 == 当前真实分类
#logging.info("%d" % obj_cls)
box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1)
pred_histro[int(obj_pred)] += 1
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thresh"]:
correct += 1
correct_histro[int(obj_pred)] += 1
break
if n_gt:
types = config["types"]
reverse_types = {} # 建立一个反向的types
for key in types.keys():
reverse_types[types[key]] = key
logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt)))
logging.info('mAP Histro:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / gt_histro[i])) for i in range(config["yolo"]["classes"] ) ]))
logging.info('Recall His:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / pred_histro[i])) for i in range(config["yolo"]["classes"]) ]))
logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
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)
# 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)
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(3):
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 == (3 - 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 %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))
# 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 / 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 train():
global_step = 0
is_training = True
# Load and Initialize Network
net = ModelMain(is_training)
net.train(is_training)
# Optimizer and Lr
optimizer = _get_optimizer(net)
lr_scheduler = optim.lr_scheduler.StepLR(
optimizer,
step_size=lr_decay_step, #20
gamma=lr_decay_gamma) # 0.1
# Set Data Paraller:
net = nn.DataParallel(net)
net = net.cuda()
logging.info("Net of Cuda is Done!")
# Restore pretrain model 从预训练模型中恢复
if pretrain_snapshot:
logging.info(
"Load pretrained weights from {}".format(pretrain_snapshot))
state_dic = torch.load(pretrain_snapshot)
net.load_state_dict(state_dic)
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(anchors[i], classes, (img_w, img_h)).cuda())
print('YOLO_Losses: \n', yolo_losses)
# DataLoader
train_data_loader = DATA.DataLoader(dataset=COCODataset(train_path,
(img_w, img_h),
is_training=True),
batch_size=batch_size,
shuffle=True,
pin_memory=False)
# Start the training loop
logging.info("Start training......")
for epoch in range(epochs):
for step, samples in enumerate(train_data_loader):
images, labels = samples['image'].cuda(), samples["label"].cuda()
start_time = time.time()
global_step += 1
# Forward & Backward
optimizer.zero_grad()
outputs = net(images)
losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"]
losses = [[]] * len(
losses_name) # [[]] ---> [[], [], [], [], [], [], []]
for i in range(3): # YOLO 3 scales
_loss_item = yolo_losses[i](outputs[i], labels)
for j, l in enumerate(_loss_item):
# print('j: ', j, 'l: ', l) j: index(0-6); l内容: 总loss, x, y, w, h, conf, cls
losses[j].append(l)
losses = [sum(l) for l in losses]
loss = losses[0] # losses[0]为总Loss
conf = losses[5]
loss.backward()
optimizer.step()
if step > 0 and step % 10 == 0:
_loss = loss.item()
_conf = conf.item()
duration = float(time.time() - start_time) # 总用时
example_per_second = batch_size / duration # 每个样本用时
lr = optimizer.param_groups[0]['lr']
logging.info(
"epoch [%.3d] iter = %d loss = %.2f conf = %.2f example/sec = %.3f lr = %.5f "
% (epoch, step, _loss, _conf, example_per_second, lr))
if step >= 0 and step % 1000 == 0:
# net.train(False)
_save_checkpoint(net.state_dict(), epoch, step)
# net.train(True)
lr_scheduler.step()
_save_checkpoint(net.state_dict(), 100, 9999)
logging.info("Bye~")
def evaluate(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"]:
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(
YOLOLoss(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=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=16,
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():
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, 80, conf_thres=0.2)
# 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 evaluate(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: {}".format(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(
YOLOLoss(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=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=False)
# Coco Prepare.
index2category = json.load(open("coco_index2category.json"))
# Start the eval loop
logging.info("Start eval.")
coco_results = []
coco_img_ids = set([])
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
image_paths, origin_sizes = samples["image_path"], samples[
"origin_size"]
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=0.01,
nms_thres=0.45)
for idx, detections in enumerate(batch_detections):
image_id = int(os.path.basename(image_paths[idx])[-16:-4])
coco_img_ids.add(image_id)
if detections is not None:
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
x1 = x1 / config["img_w"] * origin_size[0]
x2 = x2 / config["img_w"] * origin_size[0]
y1 = y1 / config["img_h"] * origin_size[1]
y2 = y2 / config["img_h"] * origin_size[1]
w = x2 - x1
h = y2 - y1
coco_results.append({
"image_id":
image_id,
"category_id":
index2category[str(int(cls_pred.item()))],
"bbox": (float(x1), float(y1), float(w), float(h)),
"score":
float(conf),
})
logging.info("Now {}/{}".format(step, len(dataloader)))
save_results_path = "coco_results.json"
with open(save_results_path, "w") as f:
json.dump(coco_results,
f,
sort_keys=True,
indent=4,
separators=(',', ':'))
logging.info("Save coco format results to {}".format(save_results_path))
# COCO api
logging.info("Using coco-evaluate tools to evaluate.")
cocoGt = COCO(config["annotation_path"])
cocoDt = cocoGt.loadRes(save_results_path)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.params.imgIds = list(coco_img_ids) # real imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
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 evaluate(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"]:
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(YOLOLoss(config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(dataset=COCODataset(config["test_path"], config["img_w"]),
batch_size=config["batch_size"],
shuffle=False, num_workers=8, pin_memory=False)
# Start the eval loop
#logging.info("Start eval.")
n_gt = 0
correct = 0
#logging.debug('%s' % str(dataloader))
gt_histro={}
pred_histro = {}
correct_histro = {}
for i in range(config["yolo"]["classes"]):
gt_histro[i] = 1
pred_histro[i] = 1
correct_histro[i] = 0
# images 是一个batch里的全部图片，labels是一个batch里面的全部标签
for step, (images, labels) in enumerate(dataloader):
labels = labels.cuda()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
# 把三个尺度上的预测结果在第1维度(第0维度是batch里的照片，第1维度是一张照片里面的各个预测框，第2维度是各个预测数值)上拼接起来
output = torch.cat(output_list, dim=1)
#logging.info('%s' % str(output.shape))
# 进行NMS抑制
#output = non_max_suppression(prediction=output, num_classes=config["yolo"]["classes"], conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"])
output = class_nms(prediction=output, num_classes=config["yolo"]["classes"],conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"])
# calculate
for sample_i in range(labels.size(0)):
# 计算所有的预测数量
sample_pred = output[sample_i]
if sample_pred is not None:
#logging.debug(sample_pred.shape)
for i in range(sample_pred.shape[0]):
pred_histro[int(sample_pred[i,6])] += 1
# Get labels for sample where width is not zero (dummies)
target_sample = labels[sample_i, labels[sample_i, :, 3] != 0]
# Ground truth的 分类编号obj_cls、相对中心x、相对中心y、相对宽w、相对高h
n_gt=0
correct=0
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
# 转化为输入像素尺寸的 左上角像素tx1 ty1，右下角像素tx2 ty2
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)
# 计算ground truth数量，用于统计信息
n_gt += 1
gt_histro[int(obj_cls)] += 1
# 转化为 shape(1,4)的tensor，用来计算IoU
box_gt = torch.cat([coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]]).view(1, -1)
# logging.info('%s' % str(box_gt.shape))
sample_pred = output[sample_i]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
# 对于每一个ground truth，遍历预测结果
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[sample_pred[:, 6] == obj_cls]: # 如果当前预测分类 == 当前真实分类
#logging.info("%d" % obj_cls)
box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1)
#pred_histro[int(obj_pred)] += 1
iou = bbox_iou(box_pred, box_gt)
#if iou >= config["iou_thres"] and obj_conf >= config["obj_thresh"]:
if iou >= config["iou_thresh"]:
correct += 1
correct_histro[int(obj_pred)] += 1
break
#logging.debug("----------------")
#logging.debug(correct_histro[4])
#logging.debug(pred_histro[4])
#logging.debug(gt_histro[4])
if n_gt:
types = config["types"]
reverse_types = {} # 建立一个反向的types
for key in types.keys():
reverse_types[types[key]] = key
#logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt)))
logging.info('Precision:%s' % str([reverse_types[i] +':'+ str(int(100 * correct_histro[i] / pred_histro[i])) for i in range(config["yolo"]["classes"]) ]))
logging.info('Recall :%s' % str([reverse_types[i] +':'+ str(int(100 * correct_histro[i] / gt_histro[i])) for i in range(config["yolo"]["classes"])]))
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 test():
is_traning = False # 不训练，测试
# Load and initialize network
net = ModelMain(is_training=is_traning)
net.train(is_traning)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if trained_model_dir:
logging.info("load checkpoint from {}".format(trained_model_dir))
state_dict = torch.load(trained_model_dir)
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(YOLOLoss(anchors[i], yolo_class_num,
(img_h, img_w)))
# prepare the images path
images_name = os.listdir("./images/")
images_path = [os.path.join("./images/", name) for name in images_name]
print('images_name：', images_name)
print('images_path：', len(images_path), images_path)
if len(images_path) == 0:
raise Exception("no image found in {}".format("./images/"))
# Start inference
batch_size = 16
for step in range(0, len(images_path),
batch_size): # range(0, 4, 16) step = 0, 4, 8, 12
logging.info('Batch_size:{}'.format(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)
# cv2.imshow('Image', image)
# cv2.waitKey(0)
logging.info(" √ Successfully Processed! √")
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) # 预处理完毕，加入原始图片组
# 进一步将图片处理为网络可以接受的数据类型（resize、归一化等）
image = cv2.resize(image, (img_h, img_w),
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()
logging.info("\nImages Convert to Tensor of CUDA Done!")
# 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(prediction=output,
num_classes=yolo_class_num,
conf_thres=0.5)
logging.info("\nNet Detection Done!\n")
# write result images: Draw BBox
classes = open(classes_name_path,
'r').read().split("\n")[:-1] # 读取coco.names
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)
# print('Final Detections: ', detections)
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 = img_h, img_w # 416, 416
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 = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
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
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}_{}.jpg'.format(step, idx),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
logging.info("All the Test Process Succeed! Enjoy it!")
def train(config):
config["global_step"] = config.get("start_step", 0)
is_training = False if config.get("export_onnx") else True
# 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)
# 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(
YOLOLoss(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),
batch_size=config["batch_size"],
shuffle=True,
num_workers=32,
pin_memory=True)
# Start the training loop
logging.info("Start training.")
for epoch in range(config["epochs"]):
for step, samples in enumerate(dataloader):
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"]
losses = []
for _ in range(len(losses_name)):
losses.append([])
for i in range(3):
_loss_item = yolo_losses[i](outputs[i], labels)
for j, l in enumerate(_loss_item):
losses[j].append(l)
losses = [sum(l) for l in losses]
loss = losses[0]
loss.backward()
optimizer.step()
if step > 0 and step % 10 == 0:
_loss = loss.item()
duration = float(time.time() - start_time)
example_per_second = config["batch_size"] / duration
lr = optimizer.param_groups[0]['lr']
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 step > 0 and step % 1000 == 0:
# net.train(False)
# _save_checkpoint(net.state_dict(), config)
# net.train(True)
_save_checkpoint(net.state_dict(), config)
lr_scheduler.step()
# net.train(False)
_save_checkpoint(net.state_dict(), config)
# net.train(True)
logging.info("Bye~")
def train(config):
# Hyper-parameters
config["global_step"] = config.get("start_step", 0)
is_training = True
# Net & Loss & Optimizer
## Net Main
net = ModelMain(config, is_training=is_training)
net.train(is_training)
## YOLO Loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_loss = YOLOLoss(config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"]))
yolo_losses.append(yolo_loss)
## Optimizer and LR scheduler
optimizer = _get_optimizer(config, net)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=config["lr"]["decay_step"], gamma=config["lr"]["decay_gamma"])
net = nn.DataParallel(net)
net = net.cuda()
# Load checkpoint
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)
# DataLoader
dataloader = torch.utils.data.DataLoader(AIPrimeDataset(config["train_path"]),
batch_size=config["batch_size"],
shuffle=True, num_workers=16, pin_memory=False)
# Start the training
logging.info("Start training.")
for epoch in range(config["start_epoch"], config["epochs"]):
for step, (images, labels) in enumerate(dataloader):
start_time = time.time()
config["global_step"] += 1
# Forward
outputs = net(images)
# Loss
losses_name = ["total_loss", "x", "y", "w", "h", "conf", "cls"]
losses = [[]] * 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].append(l)
losses = [sum(l) for l in losses]
loss = losses[0]
# Zero & Backward & Step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Logging
if step > 0 and step % 10 == 0:
_loss = loss.item()
duration = float(time.time() - start_time)
example_per_second = config["batch_size"] / duration
lr = optimizer.param_groups[0]['lr']
logging.info(
"epoch [%.3d] iter = %d loss = %.2f example/sec = %.3f lr = %.5f " %
(epoch, step, _loss, example_per_second, lr)
)
# Things to be done for every epoch
## LR schedule
lr_scheduler.step()
## Save checkpoint
_save_checkpoint(net.state_dict(), config, epoch)
# Finish training
logging.info("QiaJiaBa~ BeiBei")