letterbox

import cv2
import xml.etree.ElementTree as ET
import pdb
import os
import math
#图片拓展，使用letterbox方式填充，补黑填充
#960*300----->960*960，补黑边填充
#300*400----->400*400,

def letterbox(old_img_path, new_img_path, old_xml_path, new_xml_path):#old_img_path new_img_path old_xml_path new_xml_path
    
    lines = os.listdir(old_img_path)
    for line in lines:
        try:
            jpg_file = line.strip()
            jpg_id = jpg_file.split('.')[0] 
            xml_file = jpg_id + '.xml'
            img = cv2.imread(old_img_path + line, -1)
            ih, iw = img.shape[:2]
            eh =  max(ih, iw)
            # scale = min(eh * 1.0/ ih, eh * 1.0/ iw)

            # nh = int(ih * scale)
            # nw = int(iw * scale)
            #image = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_CUBIC)
            top = (eh - ih) // 2
            print("top:",top)
            bottom = eh - ih - top
            left = (eh - iw) // 2
            print("left:",left)
            right = eh - iw - left
            shrink = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT)
            new_jpg_file = new_img_path + jpg_file
            cv2.imwrite(new_jpg_file, shrink)

            xml_file = os.path.join(old_xml_path, xml_file)
            updateTree = ET.parse(xml_file)

            root = updateTree.getroot()
            root.find("size").find('width').text = str(eh)
            root.find("size").find('height').text = str(eh)
            element_object = root.findall('object')
            for i in range(len(element_object)):
                xmin = element_object[i].find('bndbox').find('xmin').text
                xmax = element_object[i].find('bndbox').find('xmax').text
                ymin = element_object[i].find('bndbox').find('ymin').text
                ymax = element_object[i].find('bndbox').find('ymax').text
                element_object[i].find('bndbox').find('xmin').text = str(int(int(float(xmin)) + left))
                element_object[i].find('bndbox').find('xmax').text = str(int(int(float(xmax)) + left))
                element_object[i].find('bndbox').find('ymin').text = str(int(int(float(ymin)) + top))
                element_object[i].find('bndbox').find('ymax').text = str(int(int(float(ymax)) + top))
            new_xml_file = new_xml_path + xml_file.split('/')[-1]
            updateTree.write(new_xml_file)
        except Exception as e:
            print(e)
            continue
def create_dir_not_exist(path):
    if not os.path.exists(path):
        os.mkdir(path)

if __name__ == '__main__':

    allPath = '/home/hs/data/data_jiaxingyili/goods/smalldeal'
    oldImg = allPath + '/imgC/'
    oldXml = allPath + '/xmlC/'
    newImg = allPath + '/imgL/'
    newXml = allPath + '/xmlL/'

    create_dir_not_exist(newImg)
    create_dir_not_exist(newXml)

    letterbox(oldImg, newImg, oldXml, newXml)