如果要训练一个分类模型却没有特定的分类数据集怎么办呢?可以换一种思路,将带有该目标的图片对所有想要的目标进行画标注框然后进行截图,就能得到特定的分类数据了。这么做的目的是:带有该目标的图片可能不会少,但是带有该目标且一张图内只有唯一一个目标的图片肯定少。
拿如下这张图片举例:
假设我们要训练一个车辆分类的模型,分类小车或者摩托车,首先我们先进行标注(因为只是举个例子所以就标注一部分):
然后编写截图代码:
import cv2
import json
import os
img_path = "/home/alpha/桌面/11/img"
json_path = "/home/alpha/桌面/11/json"
save_path = "/home/alpha/桌面/11/clsf"
files = os.listdir(json_path)
for file in files:
if os.path.splitext(file)[-1] != ".json":
continue
img_file = os.path.join(img_path, os.path.splitext(file)[0] + ".jpg")
json_file = os.path.join(json_path, file)
json_data = json.load(open(json_file))
i, j = 0, 0
for shape in json_data["shapes"]:
p = shape["points"]
x1 = int(min(p[0][0], p[1][0]))
y1 = int(min(p[0][1], p[1][1]))
x2 = int(max(p[0][0], p[1][0]))
y2 = int(max(p[0][1], p[1][1]))
img = cv2.imread(img_file)
img = img[y1:y2, x1:x2, :]
img_save = ""
# 根据自己的类别
if shape["label"] == "car":
img_save = save_path + "/car/" + os.path.splitext(file)[0] + "-" + str(i) + ".jpg"
i += 1
elif shape["label"] == "moto":
img_save = save_path + "/moto/" + os.path.splitext(file)[0] + "-" + str(j) + ".jpg"
j += 1
else:
continue
print(shape["label"])
cv2.imwrite(img_save, img)
print(img_save)
查看结果: