def test(self, batch_count,j):
with torch.no_grad():
self.net.eval()
img = Image.open(TEST_IMG)
# img_ = cv2.imread(TEST_IMG)
last_boxes = self.detecter.detect(img, self.args.threshold, net=self.net)
draw = ImageDraw.Draw(img)
font = ImageFont.truetype(font="arial.ttf", size=10, encoding="utf-8")
if np.any(last_boxes):
for box in last_boxes:
xybox = box[:4].astype("i4")
text_x, text_y = list(box[:2])[0], list(box[:2])[1] - 10
text_conf = list(box[:2])[0] + 30
draw.text((text_x, text_y), cfg.COCO_DICT[int(box[5])], fill=(255, 0, 0), font=font)
draw.text((text_conf, text_y), "%.2f" % box[4], fill=(255, 0, 0), font=font)
draw.rectangle(list(xybox), outline="green", width=2)
# img.show()
if NEEDSAVE:
testpic_savedir = os.path.join(SAVE_DIR, "testpic", self.netfile_name)
utils.makedir(testpic_savedir)
testpic_savefile = os.path.join(testpic_savedir, "{0}_{1}.jpg".format(batch_count,j))
img.save(testpic_savefile)
if NEEDSHOW:
plt.clf()
plt.axis("off")
plt.imshow(img)
plt.pause(0.1)
def FDplotting(self, net):
save_dir = os.path.join(SAVE_DIR, "params")
utils.makedir(save_dir)
save_name = "{0}_param.jpg".format(self.netfile_name)
save_file = os.path.join(SAVE_DIR, save_name)
params = []
for param in net.parameters():
params.extend(param.view(-1).cpu().detach().numpy())
params = np.array(params)
histo = np.histogram(params, 10, range=(np.min(params), np.max(params)))
plt.plot(histo[1][1:], histo[0])
plt.savefig(save_file)
plt.show()
def scalarplotting(self, datalist, key):
save_dir = os.path.join(SAVE_DIR, key)
utils.makedir(save_dir)
save_name = "{0}.jpg".format(key)
save_file = os.path.join(save_dir, save_name)
values = []
for data_dict in datalist:
if data_dict:
values.append(data_dict[key])
if len(values) != 0:
plt.plot(values)
plt.savefig(save_file)
plt.show()
def copydirs(src_dir,
dst_dir,
cover=False,
timethreshold=1,
outputthreshold=1000):
count = 0 # 总复制文件计数
count_eachtime = 0 # 单位时间复制文件计数
start_time = time.time() # 开始时间
time_temp = start_time # 临时时间，用于计算单位时间
# 判断文件夹存在
if os.path.isfile(src_dir):
return
utils.makedir(dst_dir)
for file in os.listdir(src_dir):
file_src = os.path.join(src_dir, file)
file_dst = os.path.join(dst_dir, file)
if os.path.isfile(file_src):
print("file_src1", file_src)
print("file_dst1", file_dst)
if cover:
shutil.copy(file_src, file_dst)
count += 1
count_eachtime += 1
elif not os.path.exists(file_dst):
shutil.copy(file_src, file_dst)
count += 1
count_eachtime += 1
check_time = time.time()
time_inter = check_time - time_temp
if count % outputthreshold == 0 and count > 0:
print("the number of files is copied successfully:", count)
if time_inter > timethreshold:
print("number of copy in {0}s: {1}, time already used: {2}s".
format(timethreshold, count_eachtime,
int(check_time - start_time)))
time_temp = check_time
count_eachtime = 0
elif os.path.isdir(file_src):
print("file_src2", file_src)
print("file_dst2", file_dst)
copydirs(file_src, file_dst, cover)
return count
def decode(self):
datadict = torch.load(self.dictpath)
for newpath, filedata in datadict.items():
savepath = os.path.join(self.dstdir, newpath).replace("\\", "/")
# if os.path.exists(savepath):
# savepath_mtime = os.stat(savepath).st_mtime
# filepath_mtime = os.stat(filepath).st_mtime
# print('s',savepath, savepath_mtime)
# print('f',filepath, filepath_mtime)
# if savepath_mtime <= filepath_mtime:
# continue
filedir = os.path.dirname(savepath)
# print("filedir", filedir)
makedir(filedir)
print("savepath", savepath)
with open(savepath, 'wb') as f:
f.writelines(filedata)
def __init__(self, net, netfile_name, cfgfile=None):
self.net = net
self.netfile_name = netfile_name
print(cfgfile)
if cfgfile != None:
self.cfginit(cfgfile)
print(1)
print(SAVE_DIR)
utils.makedir(SAVE_DIR)
parser = argparse.ArgumentParser(description="base class for network training")
self.args = self.argparser(parser)
net_savefile = "{0}.{1}".format(self.netfile_name, NETFILE_EXTENTION)
self.save_dir = os.path.join(SAVE_DIR, "nets")
utils.makedir(self.save_dir)
self.save_path = os.path.join(self.save_dir, net_savefile)
self.savepath_epoch = os.path.join(SAVEDIR_EPOCH, net_savefile)
if os.path.exists(self.save_path) and CONTINUETRAIN:
try:
self.net.load_state_dict(torch.load(self.save_path))
print("net param load successful")
except:
self.net = torch.load(self.save_path)
print("net load successful")
else:
self.net.paraminit()
print("param initial complete")
if ISCUDA:
self.net = self.net.to(DEVICE)
if NEEDTEST:
self.detecter = Detector()
self.logdir = os.path.join(SAVE_DIR, "log")
utils.makedir(self.logdir)
self.logfile = os.path.join(self.logdir, "{0}.txt".format(self.netfile_name))
if not os.path.exists(self.logfile):
with open(self.logfile, 'w') as f:
print("%.2f %d " % (0.00, 0), end='\r', file=f)
print("logfile created")
self.optimizer = optim.Adam(self.net.parameters())
# 损失函数定义
self.conf_loss_fn = nn.BCEWithLogitsLoss() # 定义置信度损失函数
self.center_loss_fn = nn.BCEWithLogitsLoss() # 定义中心点损失函数
self.wh_loss_fn = nn.MSELoss() # 宽高损失
# self.cls_loss_fn = torch.nn.CrossEntropyLoss() # 定义交叉熵损失
self.cls_loss_fn = nn.CrossEntropyLoss()
self.detecter = Detector()
print("initial complete")
def generatedataset(pass_used=False, treat_img=True, gen_label=True):
pic_used_dict = {12: "12.txt", 24: "24.txt", 48: "48.txt"}
# 新旧标签列表
datalist_read = []
# datalist_write = []
with open(labelfile, 'r') as f:
for i, line in enumerate(f.readlines()):
# 第一行为标量行，不读
if (i > 0):
datalist_read.append(line.strip())
ioulist = []
# 定义三个样本的计数
positive_num_ = 0
negative_num_ = 0
part_num_ = 0
positive_num = 0
negative_num = 0
part_num = 0
for face_size in face_sizes:
datalist_write = [] # 注意datalist清空位置，注意计数
'''
将已生成的图片存起来，下次不再生成'''
pic_used_txt = pic_used_dict.get(face_size)
if pass_used:
pic_used_list = []
if os.path.exists(pic_used_txt):
with open(pic_used_txt) as f:
for line in f.readlines():
if line[0].isdigit():
pic_used_list.append(line.strip().split()[0])
elif os.path.exists(pic_used_txt):
os.remove(pic_used_txt)
"定义三个尺寸的文件保存路径"
# positive_savepath = os.path.join(pic_savepath, str(face_size), "positive")
# negative_savepath = os.path.join(pic_savepath, str(face_size), "negative")
# part_savepath = os.path.join(pic_savepath, str(face_size), "part")
# pic_savepathlist = [negative_savepath,positive_savepath,part_savepath]
# 定义图片保存文件夹
# pic_savedict = {
# 0: os.path.join(pic_savepath, str(face_size), "negative"),
# 1: os.path.join(pic_savepath, str(face_size), "positive"),
# 2: os.path.join(pic_savepath, str(face_size), "part")
# }
'''
定义每个尺寸图片保存文件夹和标签保存文件的dict
图片要按照置信分类，标签不用'''
save_dict = {
"pic": {
0: os.path.join(pic_savepath, str(face_size), "negative"),
1: os.path.join(pic_savepath, str(face_size), "positive"),
2: os.path.join(pic_savepath, str(face_size), "part")
},
"label":
os.path.join(label_savepath,
"label_{0}.txt".format(str(face_size)))
}
# positive part negative集中放置，这几个文件夹不用生成
# for i in [0, 1, 2]:
# utils.makedir(save_dict['pic'][i])
"遍历标签数据"
for i, strdata in enumerate(datalist_read):
"分隔每行数据"
data = strdata.strip().split()
if pass_used:
if data[0] in pic_used_list:
continue
# x,y,w,h = map(int,data[1],data[2],data[3],data[4])#这个写法不行
# x1, y1, w, h, width, height = map(int, data[1:7])#本身可以实现数据的strip()
# x1, y1, w, h, width, height = map(int, map(str.strip,data[1:7]))#不以逗号分隔了，改为空格
x1, y1, x2, y2, w, h, width, height = map(int, data[1:9])
"图片过滤"
# if (x1 < 0 or y1 < 0 or w < 0 or h < 0 or max(w, h) < 40):
# continue
"判断是否有小于0的元素，框的大小是否比40小，框是否在图片内"
if any([d < 0 for d in map(int, data[1:9])
]) or min(w, h) < 40 or max(w, h) > min(width, height):
# print("false data: ", data)
# 输出错误数据
continue
'''
补方框'''
# "判断框是否为近方框"
# if w / (w + h) > 0.6 or w / (w + h) < 0.3:
# pass
# else:
# continue
# x1, y1, x2, y2, w, h, width, height = map(int, data[1:9])
"计算右下角坐标x1,y1和中心点x0,y0"
# x2, y2 = x1 + w, y1 + h
# x0, y0 = x1 + w // 2, y1 + h // 2
for j in range(len(wh_scales)):
k = 0
while k < each_pic_num:
'''
name是加载图片的名称
pic_name 是新生成图片的名称'''
name, offset, box, box_ = utils.getoffset(
data, wh_scale=wh_scales[j], size_scale=size_scale)
iou_value = utils.iou(box, box_)
# 生成图片文件名
pic_name = "%s_%d_0%1d%02d.jpg" % (
name.split('.')[0], face_size, j, k
) # 文件名中第二个"_"后的第一个0代表正或部分样本，另外生成的负样本此处为1
# 生成图片路径
"置信"
confidence = getconfidence(iou_value)
"图片保存路径"
if confidence == -1:
continue
else:
ioulist.append(iou_value)
k += 1
pic_savedir = os.path.join(pic_savepath, str(face_size))
utils.makedir(pic_savedir)
pic_savefile = os.path.join(pic_savedir, pic_name)
"标签保存路径"
label_savefile = os.path.join(save_dict['label'])
"原图片路径"
pic_file = os.path.join(picpath, name)
if treat_img:
utils.imgTreat(pic_file, pic_savefile, box_, face_size)
datalist_write.append(
utils.generateLabeldata(pic_name, confidence, offset))
# 三个样本计数
if positive_range[0] < iou_value < positive_range[1]:
positive_num_ += 1
elif part_range[0] < iou_value < part_range[1]:
part_num_ += 1
elif negative_range[0] < iou_value < negative_range[1]:
negative_num_ += 1
if i % 100 == 0:
total_num_ = positive_num_ + part_num_ + negative_num_
print("epoch", i, positive_num_, part_num_, negative_num_,
total_num_)
if all([
positive_num_ > 0, part_num_ > 0,
negative_num_ > 0, total_num_ > 0
]):
print("positive: %.3f, %.3f" %
(positive_num_ / part_num_,
positive_num_ / negative_num_))
print("ratio: %.3f, %.3f, %.3f" %
(positive_num_ / total_num_,
part_num_ / total_num_,
negative_num_ / total_num_)) # 临时查看一下数量
print("********************")
if pass_used:
with open(pic_used_txt, 'a') as pic_used_file:
print(data[0], file=pic_used_file)
if len(datalist_write) != 0 and gen_label:
utils.generateLabel(datalist_write, label_savefile, face_size)
# 三个样本计数
for i in ioulist:
if positive_range[0] < i < positive_range[1]:
positive_num += 1
elif part_range[0] < i < part_range[1]:
part_num += 1
elif negative_range[0] < i < negative_range[1]:
negative_num += 1
total_num = positive_num + part_num + negative_num
print(positive_num, part_num, negative_num, total_num)
import os import numpy as np import torch from PIL import Image, ImageDraw from tool import utils # 定义路径 # root = r"D:\datasets" "一级改动" labelfile = r"D:\datasets\datasets10261\label_10261.txt" # 读取的标签路径 "二级改动" picpath = r"D:\datasets\datasets10261\jpg" # 读取的图片路径 "一级改动" savepath = r"D:\datasets\save_10261_20190726" utils.makedir(savepath) # 图片和标签的保存根目录 pic_savepath = os.path.join(savepath, "pic") label_savepath = os.path.join(savepath, "label") "define iou threshold" positive_range = [0.65, 1] part_range = [0.4, 0.6] negative_range = [0, 0.05] size_scale = 0.9 wh_scales = [0.05, 0.1, 0.2, 0.75, 0.95] face_sizes = [12, 24, 48] classification = ['negative', 'positive', 'part'] utils.makedir(pic_savepath) utils.makedir(label_savepath)
net_2_1 原网络 [98,15]
tensor([[0, 6, 6, 0],
[0, 6, 6, 1],
[0, 6, 6, 2]], device='cuda:0')
torch.Size([3, 6])
torch.Size([98, 15])
'''
if __name__ == '__main__':
np.set_printoptions(precision=4, threshold=np.inf, suppress=True)
net_path = r"D:\PycharmProjects\yolov3_01\save\20190925\nets\yolo_03.pth"
pic_dir = r"D:\datasets\yolodatasets\datasets_20190801\datasets_resize"
save_dir = r"D:\PycharmProjects\yolov3_01\save\img"
save_dir1 = r"D:\PycharmProjects\yolov3_01\save\cv2img"
utils.makedir(save_dir)
utils.makedir(save_dir1)
# pic_name = r"pic_008.jpg"
# save_path = os.path.join(save_dir,pic_name)
torch.set_printoptions(threshold=np.inf, sci_mode=False)
detecter = Detector(net_path)
for pic_name in os.listdir(pic_dir):
pic_file = os.path.join(pic_dir, pic_name)
save_path = os.path.join(save_dir, pic_name)
# save_path1 = os.path.join(save_dir1, pic_name)
img = Image.open(pic_file)
last_boxes = detecter.detect(img, 0.6)
draw = ImageDraw.Draw(img)