图像分块代码

import cv2
import os

def judge_box_region(box, index):
    x1,y1,x2,y2 = float(box[0]),float(box[1]),float(box[2]),float(box[3])

    if index == '1':
        if x1<1280 and y1<1280:    # 左上角
            return True
        else:
            return False
    elif index == '2':
        if x2 > 768 and y1 < 1280:  # 右上角
            return True
        else:
            return False
    elif index == '3':
        if x1 < 1280 and y2 > 768:  # 左下角
            return True
        else:
            return False
    elif index == '4':
        if x2 > 768 and y2 > 768:  # 右下角
            return True
        else:
            return False

def get_index_crop_img(index, img):
    if index == '1':
        return img[:1280, :1280, :]
    elif index == '2':
        return img[:1280, 768:2048, :]
    elif index == '3':
        return img[768:2048, :1280, :]
    elif index == '4':
        return img[768:2048, 768:2048, :]

def bbox_transform(index, box):
    t_box = box.copy()
    x1,y1,x2,y2 = t_box[0],t_box[1],t_box[2],t_box[3]
    area = (x2-x1)*(y2-y1)
    #print('area: ', area)
    if index == '1':
        if x2<=1280 and y2<=1280:
            pass
        elif x2>1280 and y2<=1280: # 右边的情况
            t_box[2] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x2>1280 and y2>1280: # 右下边的情况
            t_box[2] = 1280
            t_box[3] = 1280

            crop_area = (t_box[2]-t_box[0])*(t_box[3]-t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x2<=1280 and y2>1280: # 下边的情况
            t_box[3] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        return t_box

    elif index == '2':
        if x1>=768 and y2<=1280:
            pass
        elif x1<768 and y2<=1280:  # 左边的情况
            t_box[0] = 768

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x1<768 and y2>1280:  # 左下的情况
            t_box[0] = 768
            t_box[3] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x1>=768 and y2>1280:  # 下边的情况
            t_box[3] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        t_box[0] = t_box[0] - 768 + 1
        t_box[2] = t_box[2] - 768 + 1
        return t_box

    elif index == '3':
        if x2<=1280 and y1>=768:
            pass
        elif x2<=1280 and y1<768:  # 上边的情况
            t_box[1] = 768

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x2>1280 and y1<768:  # 右上的情况
            t_box[1] = 768
            t_box[2] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x2>1280 and y1>=768:  # 右边的情况
            t_box[2] = 1280

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        t_box[1] = t_box[1] - 768 + 1
        t_box[3] = t_box[3] - 768 + 1

        return t_box

    elif index == '4':
        if x1 >= 768 and y1 >= 768:
            pass
        elif x1 >= 768 and y1 < 768:  # 上边的情况
            t_box[1] = 768

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        elif x1 < 768 and y1 < 768:  # 左上的情况
            t_box[0] = 768
            t_box[1] = 768

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False
        elif x1 < 768 and y1 >= 768:  # 左边的情况
            t_box[0] = 768

            crop_area = (t_box[2] - t_box[0]) * (t_box[3] - t_box[1])
            IoU = crop_area / area
            if IoU < 0.5:
                return False

        t_box[0] = t_box[0] - 768 + 1
        t_box[2] = t_box[2] - 768 + 1
        t_box[1] = t_box[1] - 768 + 1
        t_box[3] = t_box[3] - 768 + 1
        return t_box

if __name__ == '__main__':
    train_txt = open('train.txt', 'r')
    origin_img_path = r'F:\jittor_ttk\images'
    empty_save_img_path = r'F:\jittor_ttk\crop_data\empty'
    nonempty_save_img_path = r'F:\jittor_ttk\crop_data\nonempty'
    new_img_txt = open('new_train.txt','w')

    if not os.path.exists(empty_save_img_path):
        os.mkdir(empty_save_img_path)
    if not os.path.exists(nonempty_save_img_path):
        os.mkdir(nonempty_save_img_path)

    img_id = 0
    for num, line in enumerate(train_txt):
        # if num==100:
        #     break
        lines = str(line).strip().split(' ')

        img_name = lines[0]
        print(os.path.join(origin_img_path, str(img_name).strip()))
        img = cv2.imread(os.path.join(origin_img_path, str(img_name).strip()))
        #box_list = []
        box_num = (len(lines) - 1) // 5

        region_box_dict = {'1': [], '2': [], '3': [], '4': []}

        all_box = []
        for i in range(box_num):
            x1 = float(lines[i * 5 + 1])
            y1 = float(lines[i * 5 + 2])
            x2 = float(lines[i * 5 + 3])
            y2 = float(lines[i * 5 + 4])
            category = int(float(lines[i * 5 + 5]))

            box_list = [x1, y1, x2, y2, category]
            all_box.append(box_list)

        # region_name ——> 区域代号
        for region_name in region_box_dict.keys():
            # 循环每一个box
            for box in all_box:
                flag = judge_box_region(box, region_name)
                if flag:
                    region_box_dict[region_name].append(box)
            #print(region_box_dict)
            crop_img = get_index_crop_img(region_name, img)
            crop_img_name = '{}.jpg'.format(img_id)
            #cv2.imshow('img', crop_img)
            #cv2.waitKey(0)
            if len(region_box_dict[region_name]) == 0:
                cv2.imwrite(os.path.join(empty_save_img_path, crop_img_name), crop_img)
                continue

            cv2.imwrite(os.path.join(nonempty_save_img_path, crop_img_name), crop_img)
            img_id += 1

            line = '{}'.format(crop_img_name)
            for box in region_box_dict[region_name]:
                #print('before transform:', box)
                transform_box = bbox_transform(region_name, box)
                if not transform_box:
                    continue

                #print('after transform:', transform_box)
                x1 = float(transform_box[0])
                y1 = float(transform_box[1])
                x2 = float(transform_box[2])
                y2 = float(transform_box[3])
                category = transform_box[4]

                line += ' {} {} {} {} {}'.format(x1,y1,x2,y2,category)

            if line == '{}'.format(crop_img_name):
                continue
            new_img_txt.write('{}\n'.format(line))

    new_img_txt.close()