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)
Python 代码优化:根据坐标文件截取图像并命名

原文地址: http://www.cveoy.top/t/topic/oawd 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录