当发现摄像头采集的数据集背景区域太多了,目标只占一小部分;或者只要其中几个类别了,其余的不参与识别训练,此时这些废弃目标也将作为背景区域;又或者你想要做特定目标的拼图处理,比如我最近就在研究特定目标放在一张图的特定位置进行训练与识别,特定区域学习特定目标从而提高减小误检率…
以上的问题都可以用截图处理,这篇以最简单的根据特定类别进行截图,针对更复杂的情况的升级版本为按照一定比例外扩的部件截图。
直接上代码:
from copy import deepcopy
import cv2
import json
import os
def main(img_path, json_path, save_path, label, extend):
files = os.listdir(json_path)
for file in files:
if os.path.splitext(file)[-1] != ".json":
continue
# 要是还有非jpg、png、jpeg类型的其他格式图片,再自行添加
img_file = os.path.join(img_path, os.path.splitext(file)[0] + ".jpg")
if not os.path.exists(img_file):
img_file = os.path.join(img_path, os.path.splitext(file)[0] + ".png")
if not os.path.exists(img_file):
img_file = os.path.join(img_path, os.path.splitext(file)[0] + ".jpeg")
json_file = os.path.join(json_path, file)
json_data = json.load(open(json_file))
img_h = json_data["imageHeight"]
img_w = json_data["imageWidth"]
i = 0
for shape in json_data["shapes"]:
json_data_1 = deepcopy(json_data)
if shape["label"] in label:
img_save = os.path.join(save_path, os.path.split(img_file)[-1])
json_save = save_path + "/" + file
if os.path.exists(json_save):
json_save = save_path + "/" + str(i) + file
img_save = save_path + "/" + str(i) + os.path.split(img_file)[-1]
json_data_1["imagePath"] = str(i) + os.path.split(img_file)[-1]
i += 1
p = shape["points"]
# 外扩指定的范围
x1 = int(min(p[0][0], p[1][0])) - extend
y1 = int(min(p[0][1], p[1][1])) - extend
x2 = int(max(p[0][0], p[1][0])) + extend
y2 = int(max(p[0][1], p[1][1])) + extend
if x1 < 0:
x1 = 0
if y1 < 0:
y1 = 0
if x2 > img_w:
x2 = img_w
if y2 > img_h:
y2 = img_h
print(img_save)
print(x1, y1, x2, y2)
bbox = []
for shape1 in json_data_1["shapes"]:
# 截图的区域a内是否有别的目标框的一半以上的区域落入截图区a域内(中心点判断)有则保留作为此图(截的新图)的目标框
if shape1["label"] in label or shape1["label"] in in_label:
m_p = shape1["points"]
m_x1 = int(min(m_p[0][0], m_p[1][0]))
m_y1 = int(min(m_p[0][1], m_p[1][1]))
m_x2 = int(max(m_p[0][0], m_p[1][0]))
m_y2 = int(max(m_p[0][1], m_p[1][1]))
if not (x1 < (m_x1 + m_x2)/2 < x2 and y1 < (m_y1 + m_y2)/2 <y2):
continue
print(m_x1, m_y1, m_x2, m_y2)
bbox.append(shape1)
img = cv2.imread(img_file)
img = img[y1:y2, x1:x2, :]
json_data_1["shapes"] = []
for b_shape in bbox:
b_p = b_shape["points"]
b_p[0][0] = b_p[0][0] - x1
b_p[0][1] = b_p[0][1] - y1
b_p[1][0] = b_p[1][0] - x1
b_p[1][1] = b_p[1][1] - y1
if b_p[0][0] < 0:
b_p[0][0] = 0
if b_p[0][1] < 0:
b_p[0][1] = 0
if b_p[1][0] > x2:
b_p[1][0] = x2
if b_p[1][1] > y2:
b_p[1][1] = y2
json_data_1["shapes"].append(b_shape)
json_data_1["imageHeight"] = y2 - y1
json_data_1["imageWidth"] = x2 -x1
json.dump(json_data_1, open(json_save, "w"), ensure_ascii=False, indent=2)
cv2.imwrite(img_save, img)
if __name__ == "__main__":
img_path = "/data/wearmask/images"
json_path = "/data/wearmask/json"
save_path = "/data/wearmask/save"
label = ["head", "hat", "workhat", "helmet"] # 根据label的目标框进行截图
in_label = ["wearmask", "eye", "glasses"] # 要保留的目标框
extend = 5
main(img_path, json_path, save_path, label, extend)
先放一张原图:
这里的例图选取的不太好,背景较少没必要使用截图,但是只是为了看代码的效果。
运行后这张图:
代码的外扩范围可以自己调,这里我是外扩5个pt。截图后的左、右、下边都外扩了上边没有外扩,可以看原图,那是因为本来就外扩不了了所以进行了限制,限制了最小不能为负值,最大不能超过图片的大小,否则在运用数据集的时候会出现错误。