Python裁剪ROI區域樣本作爲訓練樣本

裁剪目標任務

 1. 裁出label2的區域塊出來

 2. 重新定義1在2中的座標

import os
import xml.etree.ElementTree as ET
from md import read_image, write_image
import cv2
from lxml import etree


def get_bbox(xml_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    object_set = root.findall('object')
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    blood_bbox_loc = []  ####

    # find object in xml files
    for object in object_set:
        obj_name = object.find('name').text
        bbox = object.find('bndbox')
        x1 = int(bbox.find('xmin').text)  # -1 #-1是因爲程序是按0作爲起始位置的
        y1 = int(bbox.find('ymin').text)  # -1
        x2 = int(bbox.find('xmax').text)  # -1
        y2 = int(bbox.find('ymax').text)  # -1
        if obj_name == '1':
            blood_bbox_loc.append([x1, y1, x2, y2])
        elif obj_name == '2':
            region_bbox_loc = [x1, y1, x2, y2]
    return region_bbox_loc, blood_bbox_loc


def crop_and_map(image, region_bbox, blood_bbox):
    ans = image[:, region_bbox[1]:region_bbox[3], region_bbox[0]:region_bbox[2]]

    new_blood_bbox = []
    for b in blood_bbox:
        # new_blood_bbox = [b[i] - region_bbox[i] for i in range(len(b))]
        b[0] = b[0] - region_bbox[0]
        b[1] = b[1] - region_bbox[1]
        b[2] = b[2] - region_bbox[0]
        b[3] = b[3] - region_bbox[1]
        new_blood_bbox.append(b)
    print(new_blood_bbox)  ####
    return ans, new_blood_bbox


def tran2xml(new_path, img_name, img, locs):
    # write to  xml
    # foldername = new_path.split('/')[-2]

    c, w, h = img.shape
    xmlname = os.path.join(new_path, img_name + '_region.xml')
    # xmlname = os.path.join(new_path, 'new', img_name + '_region.xml')

    root = etree.Element('annotation')
    etree.SubElement(root, 'filename').text = img_name + '.jpg'
    etree.SubElement(root, 'path').text = new_path + '/' + img_name + '.jpg'

    for i in range(len(locs)):
        loc = locs[i]
        xmin = loc[0]
        xmax = loc[2]
        ymin = loc[1]
        ymax = loc[3]

        obj = etree.SubElement(root, 'object')
        bbox = etree.SubElement(obj, 'bndbox')
        etree.SubElement(bbox, 'xmin').text = str(xmin)
        etree.SubElement(bbox, 'ymin').text = str(ymin)
        etree.SubElement(bbox, 'xmax').text = str(xmax)
        etree.SubElement(bbox, 'ymax').text = str(ymax)
        etree.SubElement(obj, 'difficult').text = str(0)
        etree.SubElement(obj, 'pose').text = 'Left'
        etree.SubElement(obj, 'name').text = '1'
        etree.SubElement(obj, 'truncated').text = str(1)

    # etree.SubElement(root,'folder').text=

    etree.SubElement(root, 'segmented').text = str(0)
    size = etree.SubElement(root, 'size')
    etree.SubElement(size, 'width').text = str(w)
    etree.SubElement(size, 'height').text = str(h)
    etree.SubElement(size, 'depth').text = str(c)
    tree = etree.ElementTree(root)
    tree.write(xmlname, encoding='utf-8', pretty_print=True)


if __name__ == '__main__':
    path = r'/data/dataset/xxx/naru_region_data'

    # img = read_image(os.path.join(path, '1456386107IMG-0002-00031.dcm'))
    file_name_list = sorted(
        set([file.split('.')[0] for file in os.listdir(path) if file.split('.')[-1] == 'xml']))  #####
    for file_name in file_name_list:
        image = read_image(os.path.join(path, '.'.join([file_name, 'dcm']))).to_numpy()
        region_bbox, blood_region = get_bbox(os.path.join(path, '.'.join([file_name, 'xml'])))
        ans, new_blood_bbox = crop_and_map(image, region_bbox, blood_region)

        tran2xml(path, file_name, ans, new_blood_bbox)  ####

        cv2.imwrite(os.path.join(path, '.'.join([file_name + "_region", 'jpg'])), ans.squeeze() / ans.max() * 255.0)

目前是一個半自動化的腳本，只能執行同一文件夾下的文件，便於檢查人工標註的數據是否有問題，如果xml信息缺損是無法work的，而這樣的情況通常都比較常見