def transform_image(image_bytes):
'''
字节图片流预处理
:param image_bytes: 图处的字符流
:return:
'''
my_transforms = transforms.Compose([
transforms.Resize(255), # 图片比例不变
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
image = Image.open(io.BytesIO(image_bytes)) # 读取2进制数据
if image.mode != "RGB":
__s = "input file does not RGB image..."
flog.error(' %s', __s)
raise ValueError(__s)
# 前面增加一维
__img = my_transforms(image)
# 显示测试
# import matplotlib.pyplot as plt
# plt.imshow(image)
# plt.show()
return __img.unsqueeze(0).to(device)
def handle_keypoints(self, coco_ann):
'''
:param coco_ann: 这里只会是一维标签
:return:
'''
k_ = np.array(coco_ann['keypoints'])
'''
widerface 不需要加 mask
这个是标注是否可用的的定义
# 如果关键点在物体segment内，则认为可见.
# v=0 表示这个关键点没有标注（这种情况下x=y=v=0）
# v=1 表示这个关键点标注了但是不可见(被遮挡了）
# v=2 表示这个关键点标注了同时也可见
'''
inds = np.arange(2, len(k_), 3) # 每隔两个选出 [ 2 5 8 11 14]
mask = np.ones(len(coco_ann['keypoints']), dtype=np.float32)
mask[inds] = 1
_t = k_[inds] != 2
if _t.any() and _t.sum() != len(_t):
# 有 且 不是全部的进入
_s = '标签不全为2 %s' % k_
flog.error(_s)
# raise Exception(_s)
inds = np.arange(2, len(k_), 3) # 每隔两个选出
ones = np.ones(len(coco_ann['keypoints']), dtype=np.float32)
ones[inds] = 0
nonzero = np.nonzero(ones) # 取非零索引
keypoint_np = k_[nonzero] # 选出10个数
if np.all(keypoint_np == 0):
flog.warning('出现全0 keypoints %s' % coco_ann)
return None
return keypoint_np
def f_prod_vodeo(cap,
data_transform,
model,
labels_lsit,
device,
is_keeep=False):
fps = 0.0
count = 0
num_out = 0
while True:
start_time = time.time()
'''---------------数据加载及处理--------------'''
ref, img_np = cap.read() # 读取某一帧 ref是否成功
if not ref:
if num_out >= 3:
raise Exception('摄像头无法读取~~~')
num_out += 1
flog.error('摄像头无法读取')
continue
size_ts = torch.tensor(img_np.shape[:2][::-1], device=device)
p_boxes_ltrb, p_scores_float, plabels_text = model_out4one(
model,
img_np,
data_transform,
size_ts=size_ts,
labels_lsit=labels_lsit,
is_keeep=is_keeep,
target=None)
if p_scores_float is not None and len(p_scores_float) > 0:
# 有目标画OD
img_np = f_show_od_np4cv(img_np,
p_boxes_ltrb,
p_scores_float,
plabels_text,
is_showing=False)
pass
# print("fps= %.2f" % (fps))
count += 1
img_np = cv2.putText(img_np, "fps= %.2f count=%s" % (fps, count),
(0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0),
2)
# 极小数
fps = (fps + (1. / max(sys.float_info.min,
time.time() - start_time))) / 2
cv2.imshow("video", img_np)
c = cv2.waitKey(1) & 0xff # 输入esc退出
if c == 27:
cap.release()
break
def copy_img(file_path, is_copy, path_copy_dst, path_img_src, t_names):
if is_copy:
if not os.path.exists(path_copy_dst):
os.makedirs(path_copy_dst)
for name in t_names:
_file_dst = os.path.join(path_copy_dst, name)
if os.path.exists(_file_dst):
_s = '文件有重名 %s' % _file_dst
# raise Exception(_s)
flog.error(_s)
else:
_file_src = os.path.join(path_img_src, file_path)
shutil.copy(_file_src, _file_dst)
def show_bbox_keypoints4pil(img_pil, bboxs, keypoints, scores=None):
'''
:param img_np: tensor 或 img_pil
:param boxs: np l, t, r, b
:return:
'''
if scores is None:
flog.error('无分数 %s', scores)
return
pil_copy = img_pil.copy()
draw = ImageDraw.Draw(pil_copy)
cw = 3 # 到左上角?
for bbox, k, s in zip(bboxs, keypoints, scores):
if isinstance(bboxs, np.ndarray):
l, t, r, b = bbox.astype(np.int)
elif isinstance(bboxs, torch.Tensor):
l, t, r, b = bbox.type(torch.int)
else:
raise Exception('类型错误', type(bboxs))
draw.line([(l, t), (l, b), (r, b), (r, t), (l, t)],
width=4,
fill=COLORS_ImageDraw[random.randint(
0,
len(COLORS_ImageDraw) - 1)])
# draw.chord((k[0] - cw, k[1] - cw, k[0] + cw, k[1] + cw), 0, 360, fill=(255, 0, 0), outline=(0, 255, 0))
# draw.chord((k[2] - cw, k[3] - cw, k[2] + cw, k[3] + cw), 0, 360, fill=(255, 0, 0), outline=(0, 255, 0))
# draw.chord((k[4] - cw, k[5] - cw, k[4] + cw, k[5] + cw), 0, 360, fill=(0, 0, 255), outline=(0, 255, 0))
# draw.chord((k[6] - cw, k[7] - cw, k[6] + cw, k[7] + cw), 0, 360, fill=(255, 0, 0), outline=(0, 255, 0))
# draw.chord((k[8] - cw, k[9] - cw, k[8] + cw, k[9] + cw), 0, 360, fill=(255, 0, 0), outline=(0, 255, 0))
for x, y in zip(k[::2], k[1::2]):
draw.chord((x - cw, y - cw, x + cw, y + cw),
0,
360,
fill=(255, 0, 0),
outline=(0, 255, 0))
# draw.text((l, t), "hello", (0, 255, 0))
# draw.point([(20, 20), (25, 25), (50, 50), (30, 30)], (0, 255, 0))
# ltrb 角度顺时间 框色 填充色
# draw.chord((0, 0, 3, 3), 0, 360, fill=(255, 0, 0), outline=(0, 255, 0))
pil_copy.show()
def __init__(self, path_data_root, path_file_txt, transforms=None,
bbox2one=False, isdebug=False, is_mosaic=False, cfg=None):
'''
:param path_data_root: voc数据集的根目录
:param path_file_txt: 提前准备好的xml及jpg对应的文件名
:param transforms: 自定义的 transforms 支持 boxes 一起处理
:param bbox2one: bbox是否需要统一化 bbox2one
'''
self.path_data_root = path_data_root
self.transforms = transforms
self.bbox2one = bbox2one
self.isdebug = isdebug
self.is_mosaic = is_mosaic
self.cfg = cfg
path_txt = os.path.join(path_data_root, path_file_txt)
_path_xml = os.path.join(path_data_root, 'Annotations')
with open(
path_txt) as read: # {FileNotFoundError}[Errno 2] No such file or directory: '/home/bak3t/bakls299g/AI/datas/VOC2012/trainval/train.txt'
# 读每一行加上路径和扩展名---完整路径
self.xml_list = [os.path.join(_path_xml, line.strip() + ".xml")
for line in read.readlines()]
try:
# {"类别1": 1, "类别2":2}
path_json_class = os.path.abspath(os.path.join(path_data_root, ".."))
json_file = open(os.path.join(path_json_class, 'classes_ids_voc.json'), 'r')
self.class_to_ids = json.load(json_file)
self.ids_to_class = dict((val, key) for key, val in self.class_to_ids.items())
file_json_ids_class = os.path.join(path_json_class, 'ids_classes_voc.json')
if not os.path.exists(file_json_ids_class):
json_str = json.dumps(self.ids_to_class, indent=4)
file_json_ids_class = os.path.join(path_json_class, 'ids_classes_voc.json')
with open(file_json_ids_class, 'w') as json_file:
json_file.write(json_str)
except Exception as e:
flog.error(e)
exit(-1)
def mgpu_process0_init(args, cfg, loader_train, loader_val_coco, model,
device):
'''支持 add_graph'''
# 主进程任务
flog.info('多GPU参数: %s' % args)
if not os.path.exists(cfg.PATH_SAVE_WEIGHT):
try:
os.makedirs(cfg.PATH_SAVE_WEIGHT)
except Exception as e:
flog.error(' %s %s', cfg.PATH_SAVE_WEIGHT, e)
# tensorboard --logdir=runs_widerface --host=192.168.0.199
# tensorboard --logdir=runs_voc --host=192.168.0.199
# print('"tensorboard --logdir=runs_raccoon200 --host=192.168.0.199", view at http://192.168.0.199:6006/'
# % cfg.PATH_TENSORBOARD)
import time
c_time = time.strftime('%Y-%m-%d_%H_%M_%S', time.localtime(time.time()))
path = os.path.join(cfg.PATH_PROJECT_ROOT, 'log', cfg.PATH_TENSORBOARD,
c_time)
os.makedirs(path, exist_ok=True)
# img_ = None
# if os.path.exists(path):
# if cfg.DEL_TB and cfg.PATH_HOST == '':
# # os.remove(path) # 删除空文件夹 shutil.rmtree(path, ignore_errors=True)
# os.system("rm -rf %s" % path) # Linux下调用bash命令
# img_ = torch.zeros((1, 3, *cfg.IMAGE_SIZE), device=device) # 删除后需要
# import time
# while os.path.exists(path):
# time.sleep(1)
# else:
# flog.error('tb_writer 删除成功: %s', path)
# pass
# else:
# img_ = torch.zeros((1, 3, *cfg.IMAGE_SIZE), device=device) # 不存在时需要
tb_writer = SummaryWriter(path)
# if img_ is not None:
# tb_writer.add_graph(model, img_)
return tb_writer
def set_tail(self):
self.cfg.PATH_SAVE_WEIGHT = self.cfg.PATH_HOST + '/AI/weights/feadre'
self.cfg.FILE_FIT_WEIGHT = os.path.join(self.cfg.PATH_SAVE_WEIGHT,
self.cfg.FILE_NAME_WEIGHT)
# json_file = open(os.path.join(self.cfg.PATH_DATA_ROOT, 'ids_classes.json'), 'r', encoding='utf-8')
# self.cfg.IDS_CLASSES = json.load(json_file, encoding='utf-8') # json key是字符
if self.cfg.IS_FMAP_EVAL:
f_recover_gt(self.cfg.PATH_EVAL_INFO + '/gt_info')
# device = torch.device("cpu")
self.cfg.DATA_NUM_WORKERS = min(
[os.cpu_count(), self.cfg.DATA_NUM_WORKERS])
# self.cfg.DATA_NUM_WORKERS = 0 # 这里设置数据线程
# 检查保存权重文件夹是否存在，不存在则创建
if not os.path.exists(self.cfg.PATH_SAVE_WEIGHT):
try:
os.makedirs(self.cfg.PATH_SAVE_WEIGHT)
except Exception as e:
flog.error(' %s %s', self.cfg.PATH_SAVE_WEIGHT, e)
def load_anns(self, index, img_wh):
'''
ltwh --> ltrb
:param index:
:return:
bboxs: np(num_anns, 4)
labels: np(num_anns)
'''
# annotation_ids = self.coco.getAnnIds(self.image_ids[index], iscrowd=False)
annotation_ids = self.coco_obj.getAnnIds(self.ids_img[index]) # ann的id
# anns is num_anns x 4, (x1, x2, y1, y2)
bboxs_np = np.zeros((0, 4), dtype=np.float32) # np创建 空数组 高级
labels = []
keypoints_np = []
# skip the image without annoations
if len(annotation_ids) == 0:
return None
coco_anns = self.coco_obj.loadAnns(annotation_ids)
for a in coco_anns:
x, y, box_w, box_h = a['bbox'] # ltwh
# 得 ltrb
x1 = max(0, x) # 修正lt最小为0 左上必须在图中
y1 = max(0, y)
x2 = min(img_wh[0] - 1, x1 + max(0, box_w - 1)) # 右下必须在图中
y2 = min(img_wh[1] - 1, y1 + max(0, box_h - 1))
if a['area'] > 0 and x2 >= x1 and y2 >= y1:
a['bbox'] = [x1, y1, x2, y2] # 修正 并写入ltrb
else:
flog.warning('标记框有问题 %s 跳过', a)
continue
bbox = np.zeros((1, 4), dtype=np.float32)
bbox[0, :4] = a['bbox']
if self.mode == 'keypoints':
'''
# 如果关键点在物体segment内，则认为可见.
# v=0 表示这个关键点没有标注（这种情况下x=y=v=0）
# v=1 表示这个关键点标注了但是不可见(被遮挡了）
# v=2 表示这个关键点标注了同时也可见
'''
keypoints = self.handle_keypoints(a)
if keypoints is None:
flog.warning('全0 keypoints %s 跳过')
continue
keypoints_np.append(keypoints)
# 全部通过后添加
bboxs_np = np.append(bboxs_np, bbox, axis=0)
labels.append(self.ids_old_new[a['category_id']])
# bboxs = ltwh2ltrb(bboxs) # 前面 已转
if bboxs_np.shape[0] == 0:
flog.error('这图标注 不存在 %s', coco_anns)
return None
# raise Exception('这图标注 不存在 %s', coco_anns)
# 转tensor
if self.mode == 'bbox':
return [torch.tensor(bboxs_np, dtype=torch.float), torch.tensor(labels, dtype=torch.float)]
elif self.mode == 'keypoints':
keypoints_np = np.array(keypoints_np) # list转np
# 有标注的情况下 keypoints_np 一定有
return [torch.tensor(bboxs_np, dtype=torch.float),
torch.tensor(labels, dtype=torch.float),
torch.tensor(keypoints_np, dtype=torch.float)]
def __getitem__(self, index):
'''
:param index:
:return: tensor or np.array 根据 out: 默认ts or other is np
img: h,w,c
target:
coco原装是 ltwh
dict{
image_id: int,
bboxs: ts n4 原图 ltwh -> ltrb
labels: ts n,
keypoints: ts n,10
size: wh
}
'''
# 这里生成的是原图尺寸的 target 和img_np_uint8 (375, 500, 3)
if self.is_mosaic and self.mode == 'bbox':
res = self.do_mosaic(index)
else:
res = self.open_img_tar(index)
if res is None:
flog.error('这个图片没有标注信息 id为%s', index)
return self.__getitem__(index + 1)
img, target = res
if len(target['boxes']) != len(target['labels']):
flog.warning('!!! 数据有问题 1111 %s %s %s ', target, len(target['boxes']), len(target['labels']))
'''---------------cocoAPI测试 查看图片在归一化前------------------'''
# 这个用于调试
# if self.cfg.IS_VISUAL_PRETREATMENT:
# 可视化参数 is_mosaic 这个用不起
# f_show_coco_pics(self.coco_obj, self.path_img, ids_img=[index])
if target['boxes'].shape[0] == 0:
flog.warning('数据有问题 重新加载 %s', index)
return self.__getitem__(index + 1)
if self.transform is not None:
img, target = self.transform(img, target)
# if self.is_img_np:
# # 输入 ltrb 原图
# # f_plt_show_cv(img, gboxes_ltrb=target['boxes'])
# # img, boxes, labels = self.transform(img, target['boxes'], target['labels'])
# img, target = self.transform(img, target)
# # 这里会刷新 boxes, labels
# # f_plt_show_cv(img, gboxes_ltrb=boxes)
# else:
# # 预处理输入 PIL img 和 np的target
# img, target = self.transform(img, target)
if len(target['boxes']) != len(target['labels']):
flog.warning('!!! 数据有问题 ttttttttt %s %s %s ', target, len(target['boxes']), len(target['labels']))
# target['boxes'] = torch.tensor(target['boxes'], dtype=torch.float)
# target['labels'] = torch.tensor(target['labels'], dtype=torch.int64)
target['size'] = torch.tensor(target['size'], dtype=torch.float) # 用于恢复真实尺寸
# if self.mode == 'keypoints':
# target['keypoints'] = torch.tensor(target['keypoints'], dtype=torch.float)
if target['boxes'].shape[0] == 0:
flog.debug('二次检查出错 %s', index)
return self.__getitem__(index + 1)
if len(target['boxes']) != len(target['labels']):
flog.warning('!!! 数据有问题 22222 %s %s %s ', target, len(target['boxes']), len(target['labels']))
# flog.warning('数据debug 有问题 %s %s %s ', target, len(target['boxes']), len(target['labels']))
return img, target
def load_weight(file_weight,
model,
optimizer=None,
lr_scheduler=None,
device=torch.device('cpu'),
is_mgpu=False,
ffun=None):
start_epoch = 0
# 只匹配需要的
# model_dict = model.state_dict() # 获取模型每层的参数阵
# checkpoint = torch.load(file_weight, map_location=device)
# pretrained_dict = checkpoint['model']
# dd = {}
# for k, v in pretrained_dict.items():
# if k in model_dict:
# shape_1 = model_dict[k].shape
# shape_2 = pretrained_dict[k].shape
# if shape_1 == shape_2:
# dd[k] = v
# else:
# print('shape mismatch in %s. shape_1=%s, while shape_2=%s.' % (k, shape_1, shape_2))
# model_dict.update(dd)
# model.load_state_dict(model_dict)
# flog.warning('手动加载完成:%s',file_weight)
# return start_epoch
if file_weight and os.path.exists(file_weight):
checkpoint = torch.load(file_weight, map_location=device)
'''对多gpu的k进行修复'''
if 'model' in checkpoint:
pretrained_dict_y = checkpoint['model']
else:
pretrained_dict_y = checkpoint
# 特殊处理
# if True:
# del pretrained_dict['module.head_yolov1.weight']
# del pretrained_dict['module.head_yolov1.bias']
# # del pretrained_dict['module.ClassHead.1.conv1x1.weight']
# # del pretrained_dict['module.ClassHead.1.conv1x1.bias']
# # del pretrained_dict['module.ClassHead.2.conv1x1.weight']
# # del pretrained_dict['module.ClassHead.2.conv1x1.bias']
# model.load_state_dict(pretrained_dict, strict=False)
# start_epoch = checkpoint['epoch'] + 1
# flog.error('已特殊加载 feadre 权重文件为 %s', file_weight)
# return start_epoch
''' debug '''
if ffun is not None:
pretrained_dict = ffun(pretrained_dict_y)
else:
pretrained_dict = pretrained_dict_y
dd = {}
# 多GPU处理
ss = 'module.'
for k, v in pretrained_dict.items():
if is_mgpu:
if ss not in k:
dd[ss + k] = v
else:
dd = pretrained_dict_y
break
# dd[k] = v
else:
dd[k.replace(ss, '')] = v
'''重组权重'''
# load_weights_dict = {k: v for k, v in weights_dict.items()
# if model.state_dict()[k].numel() == v.numel()}
keys_missing, keys_unexpected = model.load_state_dict(dd, strict=False)
if len(keys_missing) > 0 or len(keys_unexpected):
flog.error('missing_keys %s', keys_missing) # 这个是 model 的属性
flog.error('unexpected_keys %s', keys_unexpected) # 这个是 pth 的属性
if optimizer:
optimizer.load_state_dict(checkpoint['optimizer'])
if lr_scheduler and checkpoint['lr_scheduler']:
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
if 'epoch' in checkpoint:
start_epoch = checkpoint['epoch'] + 1
else:
start_epoch = 0
flog.warning('已加载 feadre 权重文件为 %s', file_weight)
else:
# raise Exception(' 未加载 feadre权重文件 ')
flog.error(' 未加载 feadre权重文件 %s', file_weight)
if fis_mgpu():
path = os.path.join(tempfile.gettempdir(), "_init_weights_tmp.pt")
checkpoint_path = path
if is_main_process():
torch.save(model.state_dict(), checkpoint_path)
# import torch.distributed as dist
torch.distributed.barrier() # 多GPU阻塞
model.load_state_dict(
torch.load(checkpoint_path, map_location=device))
flog.error('已默认加载临时文件 %s', path)
return start_epoch
def cre_transform_resize4np(cfg):
if cfg.USE_BASE4NP:
flog.error('使用的是 USE_BASE4NP 模式 %s', cfg.USE_BASE4NP)
data_transform = {
"train":
Compose(
[
ConvertFromInts(), # image int8转换成float [0,256)
Resize(cfg.IMAGE_SIZE),
Normalize(cfg.PIC_MEAN, cfg.PIC_STD),
ConvertColor(current='BGR', transform='RGB'),
ToTensor(is_box_oned=False),
],
cfg)
}
else:
if cfg.KEEP_SIZE: # 不进行随机缩放 不改变图片的位置和大小
data_transform = {
"train":
Compose(
[
ConvertFromInts(), # image int8转换成float [0,256)
# ToAbsoluteCoords(), # 恢复真实尺寸
PhotometricDistort(), # 图片处理集合
# Expand(cfg.PIC_MEAN), # 放大缩小图片
# RandomSampleCrop(), # 随机剪切定位
RandomMirror(),
# ToPercentCoords(), # boxes 按原图归一化 最后ToTensor 统一归一
Resize(cfg.IMAGE_SIZE), # 定义模型输入尺寸 处理img和boxes
Normalize(cfg.PIC_MEAN, cfg.PIC_STD), # 正则化图片
ConvertColor(current='BGR', transform='RGB'),
ToTensor(is_box_oned=True
), # img 及 boxes(可选,前面已归一)归一 转tensor
],
cfg)
}
else:
data_transform = {
"train":
Compose(
[
ConvertFromInts(), # image int8转换成float [0,256)
# ToAbsoluteCoords(), # 输入已是原图不需要恢复 boxes 恢复原图尺寸
PhotometricDistort(), # 图片处理集合
Expand(cfg.PIC_MEAN), # 放大缩小图片
RandomSampleCrop(), # 随机剪切定位 keypoints
RandomMirror(),
# ToPercentCoords(), # boxes 按原图归一化 最后统一归一 最后ToTensor 统一归一
Resize(cfg.IMAGE_SIZE),
Normalize(cfg.PIC_MEAN, cfg.PIC_STD),
ConvertColor(current='BGR', transform='RGB'),
ToTensor(is_box_oned=True),
],
cfg)
}
data_transform["val"] = Compose([
Resize(cfg.IMAGE_SIZE),
Normalize(cfg.PIC_MEAN, cfg.PIC_STD),
ConvertColor(current='BGR', transform='RGB'),
ToTensor(is_box_oned=False),
], cfg)
return data_transform
