Python 代码优化:根据坐标文件截取图像并命名
import os
import json
import xml.etree.ElementTree as ET
from PIL import Image
import cv2
def read_coordinates(filepath, file_type):
'''
根据文件类型读取坐标信息
'''
if file_type == 'json':
with open(filepath, 'r') as f:
data = json.load(f)
if 'coordinates' in data:
return data['coordinates']
elif file_type == 'xml':
root = ET.parse(filepath).getroot()
coord_list = []
for obj in root.iter('object'):
coords = obj.find('bndbox')
name = obj.find('name').text
if name == 'plate':
coord_list.append([
int(coords.find('xmin').text),
int(coords.find('ymin').text),
int(coords.find('xmax').text),
int(coords.find('ymax').text),
'0'
])
elif name == 'double':
coord_list.append([
int(coords.find('xmin').text),
int(coords.find('ymin').text),
int(coords.find('xmax').text),
int(coords.find('ymax').text),
'01'
])
return coord_list
elif file_type == 'txt':
with open(filepath, 'r') as f:
coords = f.readline().strip()
coords = [int(coord) for coord in coords.split(',')[:-1]]
return coords
else:
raise ValueError('Unknown file type')
return None
def capture_image(image_path, coords, var, save_start, filename, method='cv2'):
'''
实现根据坐标截取图像
'''
if method == 'cv2':
img = cv2.imread(image_path)
xmin, ymin, xmax, ymax, suffix = coords
xmin = max(0, xmin - var)
ymin = max(0, ymin - var)
xmax = min(img.shape[1], xmax + var)
ymax = min(img.shape[0], ymax + var)
img = img[ymin:ymax, xmin:xmax]
x1,y1,x2,y2 = coords[:4]
sx1 = x1 - xmin
sx2 = x2 - xmin
sy1 = y1 - ymin
sy2 = y2 - ymin
file_suffix = filename.split('.')[-1]
save_path = f'{save_start}_{suffix}_{sx1:02d}{sy1:02d}{sx2:02d}{sy2:02d}.{file_suffix}'
cv2.imwrite(save_path, img)
print('保存截图:', save_path)
elif method == 'PIL':
img = Image.open(image_path)
xmin, ymin, xmax, ymax, suffix = coords
xmin = max(0, xmin - var)
ymin = max(0, ymin - var)
xmax = min(img.width, xmax + var)
ymax = min(img.height, ymax + var)
img = img.crop((xmin, ymin, xmax, ymax))
# img.save(save_path)
if __name__ == '__main__':
image_dir = 'F:\park\C06right'
if not os.path.isdir(image_dir):
print('无效的图片文件夹路径')
exit(1)
coord_dir = 'F:\park\C06right_xml'
if not os.path.isdir(coord_dir):
print('无效的坐标文件夹路径')
exit(1)
var = 150
method = 'cv2'
save_dir = 'F:\park\C06right_cut'
if not os.path.isdir(save_dir):
print('无效的保存文件夹路径')
exit(1)
for filename in os.listdir(image_dir):
if filename.endswith(('jpg', 'jpeg', 'png', 'JPG', 'JPEG', 'PNG')):
basename = os.path.splitext(filename)[0]
coord_path = os.path.join(coord_dir, basename + '.xml')
if not os.path.isfile(coord_path):
coord_path = os.path.join(coord_dir, basename + '.json')
if not os.path.isfile(coord_path):
coord_path = os.path.join(coord_dir, basename + '.txt')
if not os.path.isfile(coord_path):
continue
coords = read_coordinates(coord_path, os.path.splitext(coord_path)[-1][1:])
if coords is None:
continue
img_path = os.path.join(image_dir, filename)
for i, coord in enumerate(coords):
save_start = os.path.join(save_dir, f'{basename}_{coord[4]}_{i}')
capture_image(img_path, coord, var, save_start, filename, method)
原文地址: http://www.cveoy.top/t/topic/oawd 著作权归作者所有。请勿转载和采集!