这个脚本的由来是最近需要研究算法如何识别目标,这里我们用到了detection2算法模型来进行训练,由于我们训练模型开始时需要找大量的数据集,官方给出的数据集不够完善,无法训练出来比较合适的模型,整理数据集的时候发现Github上有labelme实现切图的json文件编写(这个方法很好找网上有很多有关labelme实现编写coco数据集的实例),但是没有将切图过程和写json文件coco数据集连贯起来,说到这大家就知道了这个项目如何诞生的了吧。
这个项目我打算开源出来,大家需要的话可以拉取代码:djangotest: 切片算法、向船舶识别算法送切出来的文件路径https://gitee.com/wangsenKeyL/django-test 下面我来详细讲一下实现过程。
一、切图实现
切图算法实现起来很简单,我这边切的是后缀为.tif的大图,切图过程中需要保存栅格数据的经纬度,需要用到gdal库,并添加步长解决切图过程中切割不完整的问题。
if len(data.shape) == 3:
bands, height, width = data.shape
# 切割像素
size = 1500
step = 500
y = math.ceil((height - size) // step) + 1
x = math.ceil((width - size) // step) + 1
print(x)
print(y)
for j in range(y): # 切割成小图
for i in range(x):
if j != y - 1:
if i != x - 1:
cur_image = data[bands - 1, i * step:i * step + size, j * step:j * step + size]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
cur_image = data[bands - 1, -size:, j * step:j * step + size]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
if i != x - 1:
cur_image = data[bands - 1, i * step:i * step + size, -size:]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
cur_image = data[bands - 1, -size:, -size:]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
height, width = data.shape
size = 1500
step = 500
y = math.ceil((height - size) // step) + 1
x = math.ceil((width - size) // step) + 1
print(x)
print(y)
for j in range(y): # 切割成小图
for i in range(x):
if j != y - 1:
if i != x - 1:
cur_image = data[i * step:i * step + size, j * step:j * step + size]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
cur_image = data[-size:, j * step:j * step + size]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
if i != x - 1:
cur_image = data[i * step:i * step + size, -size:]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
else:
cur_image = data[-size:, -size:]
lon = minx + xres * step * j
lat = maxy + yres * (i * step)
util.write_img(r'D:\faceWatch\temptif/{}_{}.tif'.format(i, j), proj,
lon, lat, xres, yres, cur_image) # 写数据
二、脚本实现coco数据集所需的json文件自动化生成
这里由于需要遍历指定文件夹中所有图片,我这里过程文件生成两个json文件然后拼接完成最终的json文件编写。
# 找到切图文件所在的文件夹,然后把文件夹里包含的文件(以.tif/.jpg为后缀的)名写入到json文件中
id = -1
path = "D:\\faceWatch\\temptif\\"
file_list = nt.natsorted(os.listdir(path))
print(file_list)
jsonObject = []
for filename in file_list:
id += 1
print("=======write======")
# print(filename)
dirs = path + filename
img = cv2.imread(dirs)
size = img.shape
# print(size)
height = size[0]
width = size[1]
# print(height)
# print(width)
if id <= 0:
test_dict = {
"info": {"descroption": "null", "url": "null", "version": "null", "year": 2023,
"contributor": "null",
"date_created": "2023-06-06"},
"licenses": [{"url": "null", "id": 0, "name": "null"}],
"images": [
{"license": 0, "url": "null", "file_name": filename, "height": height,
"width": width, "date_captured": "null", "id": id}
],
"type": "instances",
"annotations": [],
"categories": [
{
"supercategory": "null",
"id": 1,
"name": "ship"
}
]
},
json_str = json.dumps(test_dict)
new_dict = json.loads(json_str)
with open(r"D:\Pycharm\django-test\config\record.json", "w", encoding='utf-8') as fw:
json.dump(new_dict, fw, indent=4, ensure_ascii=False)
print("\n")
else:
test_dict_content = {"license": 0, "url": "null", "file_name": filename, "height": height,
"width": width, "date_captured": "null", "id": id}
jsonObject.append(test_dict_content)
json_str = json.dumps(jsonObject)
new_dict = json.loads(json_str)
with open(r"D:\Pycharm\django-test\config\record1.json", "w", encoding='utf-8') as fw:
json.dump(new_dict, fw, indent=4, ensure_ascii=False)
print("\n")
with open(r"D:\Pycharm\django-test\config\record.json") as fin1:
data1 = json.load(fin1)
with open(r"D:\Pycharm\django-test\config\record1.json") as fin2:
data2 = json.load(fin2)
data1[0]["images"].extend(
data2) # extend() 方法用于将一个可迭代对象中的元素逐个追加到列表中。它将迭代遍历可迭代对象并将每个元素逐个添加到列表中。使用 extend() 将 data2
# 中的元素逐个添加到 data1[0]["images"] 中,可以保持列表的扁平结构,避免额外的嵌套。补充:append方法用于将一个元素追加到列表的末尾,作为一个单独的元素。如果使用append()将data2
# 添加到data1[0]["images"]中,导致最外层多了一个列表。
json_str = json.dumps(data1)
print(json_str)
old_dict = json_str[1:-1]
print(old_dict)
new_dict = json.loads(old_dict)
print(new_dict)
print(type(new_dict))
number = random.randint(0, 10000)
json_name = "testJson_" + str(number) + ".json"
with open(r"D:/Pycharm/django-test/config/" + json_name, "w") as fin3:
json.dump(new_dict, fin3)
print(number)
实现过程如上,编写完成的json文件可以直接给mmdetection、detection2模型进行测试和训练,这里封装django框架的目的就是为了服务通信,将最终的json文件传给算法进行识别。
三、小结
这个脚本的编写只是为了实现简单的切图+写json文件的连贯,由于没有用到多线程针对大量的数据集进行并行处理,读写的效率还是蛮低的,大家如果有比较好的并行处理的办法可以分享出来,欢迎投稿我的邮箱wangsenKeyL@163.com.