如果要训练一个分类模型却没有特定的分类数据集怎么办呢?可以换一种思路,将带有该目标的图片对所有想要的目标进行画标注框然后进行截图,就能得到特定的分类数据了。这么做的目的是:带有该目标的图片可能不会少,但是带有该目标且一张图内只有唯一一个目标的图片肯定少。
拿如下这张图片举例:
假设我们要训练一个车辆分类的模型,分类小车或者摩托车,首先我们先进行标注(因为只是举个例子所以就标注一部分):
然后编写截图代码:
import cv2
import json
import osimg_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":continueimg_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, 0for 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 += 1elif shape["label"] == "moto":img_save = save_path + "/moto/" + os.path.splitext(file)[0] + "-" + str(j) + ".jpg"j += 1else:continueprint(shape["label"])cv2.imwrite(img_save, img)print(img_save)
查看结果: