plt原图像素+去四周白边和坐标轴保存图片

在复现一个floor_plan的语义分割的任务的时候想提取其中的一个子任务，也就是平面图识别中的房间识别，于是我需要把标记好的图片通过plt保存，并且维持原分辨率（像素）并且去除坐标轴和白边，看了很多教程，探索出了一段可以批量处理Cubicasa数据集房间识别的代码：

import cv2
import shutil
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from floortrans.loaders import FloorplanSVG
from torch.utils.data import DataLoader
from floortrans.plotting import discrete_cmap

discrete_cmap()

room_classes = ["Background", "Outdoor", "Wall", "Kitchen", "Living Room" ,"Bed Room", "Bath", "Entry", "Railing", "Storage", "Garage", "Undefined"] #不同的房间类型
data_folder = 'data/cubicasa140/' # 数据集文件夹
data_file = 'train.txt' # 数据集类型
savedata_folder = 'data_/pic/'
savelabel_folder = 'data_/label/' 

normal_set = FloorplanSVG(data_folder, data_file, format='txt', original_size=True) #加载数据集
data_loader = DataLoader(normal_set, batch_size=1, num_workers=0) #把加载的数据集封装起来
data_iter = iter(data_loader)

for data in tqdm(data_iter):
    folder = data['folder'][0]
    image = data['image']
    label = data['label']
    
    label_np = label.data.numpy()[0]
    plt.figure(figsize=(label_np.shape[-1],label.shape[-2]), dpi=1) # 保持figure_size的大小和原图的大小不变
    plt.gca().xaxis.set_major_locator(plt.NullLocator()) # 不显示坐标轴上的刻度
    plt.gca().yaxis.set_major_locator(plt.NullLocator()) # 不显示坐标轴上的刻度
    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0) # 子图的顶部和底部与图表的边缘对齐
    plt.margins(0,0) # 这一句代码将图表的边缘空白区域设置为0
    ax = plt.subplot(1, 1, 1)
    ax.axis('off') # 不要刻度

    n_rooms = 12
    rseg = ax.imshow(label_np[0], cmap='rooms', vmin=0, vmax=n_rooms-0.1)

    save_name = folder[1:].replace('/', '_')
    plt.savefig(savelabel_folder+save_name+'.jpg', dpi=1) # # 保持figure_size的大小和原图的大小不变
    
    src = 'data/cubicasa140' + folder + 'F1_original.png'
    dst = savedata_folder+save_name+'.jpg'
    shutil.copyfile(src, dst)
    plt.close()
    
    # print(save_name+'数据生成完成')
    

参考链接：知乎