神经网络训练中PNG转np.array异常问题及解决

我自定义了AdjustImage变换类，用于调整X光图像的旋转角度，适配神经网络训练需求。单独测试单张PNG或JPG图像时该类功能正常，但将其接入PyTorch训练流程后，PNG图像转换为np.array时出现异常，生成全黑图像，导致后续边缘检测无法执行，最终在寻找凸包的环节报错。

自定义AdjustImage变换类代码 #

class AdjustImage(object):
    THRESHOLD = 240
    def subimage(self, image, center, theta, width, height):
        if 45 < theta <= 90:
            theta = theta - 90
            width, height = height, width
    
        theta *= math.pi / 180 # 转换为弧度
        v_x = (math.cos(theta), math.sin(theta))
        v_y = (-math.sin(theta), math.cos(theta))
        s_x = center[0] - v_x[0] * (width / 2) - v_y[0] * (height / 2)
        s_y = center[1] - v_x[1] * (width / 2) - v_y[1] * (height / 2)
        mapping = np.array([[v_x[0],v_y[0], s_x], [v_x[1],v_y[1], s_y]])
        return cv2.warpAffine(image, mapping, (width, height), flags=cv2.WARP_INVERSE_MAP, borderMode=cv2.BORDER_REPLICATE)
    def __call__(self, image_source):
         # 先裁剪边缘部分 - 部分图像单边有2像素的黑边
        type_var = image_source.format
        image_source.show(title="初始图像")
        if isinstance(image_source, np.ndarray):
            print("working")
        elif image_source.format == 'JPEG':
            image_source = np.array(image_source)
            image_source = image_source[:,:,::-1].copy()
        else:
            image_source = image_source.point(lambda x: x / 256)
            image_source = image_source.convert('RGB')
            image_source.show(title="转换后图像")
            image_source = np.array(image_source)
            image_source = image_source[:,:,::-1].copy()
            
            # 该方法基于cv2.imread测试开发，这些步骤将PIL读取的PNG转换为对应格式
            # 否则会损坏图像
           
        init_crop = 5
        
        h, w = image_source.shape[:2]
        image_source = image_source[init_crop:init_crop+(h-init_crop*2), init_crop:init_crop+(w-init_crop*2)]
        
        # 添加白色边框
        image_source = cv2.copyMakeBorder(image_source, 5,5,5,5, cv2.BORDER_CONSTANT, value=(255,255,255))
        
        image_gray = cv2.cvtColor(image_source, cv2.COLOR_BGR2GRAY)
        _, image_thresh = cv2.threshold(image_gray, self.THRESHOLD, 255, cv2.THRESH_TOZERO_INV)
        
        image_thresh2 = image_thresh.copy()
        image_thresh2 = cv2.Canny(image_thresh2, 100, 100, apertureSize=3)
        
        points = cv2.findNonZero(image_thresh2)
    
        centre, dimensions, theta = cv2.minAreaRect(points)
        rect = cv2.minAreaRect(points)
    
        width = int(dimensions[0])
        height = int(dimensions[1])
    
        box = cv2.boxPoints(rect)
        box = np.int0(box)
    
        temp = image_source.copy()
        cv2.drawContours(temp, [box], 0, (255,0,0), 2)
    
        M = cv2.moments(box)    
        cx = int(M['m10']/M['m00'])
        cy = int(M['m01']/M['m00'])
    
        image_patch = self.subimage(image_source, (cx, cy), (theta+90), height, width)
    
        # 添加小边框
        image_patch = cv2.copyMakeBorder(image_patch, 1,1,1,1, cv2.BORDER_CONSTANT, value=(255,255,255))
    
        # 转换为二值图，边缘为黑色。执行边缘检测并将边缘转为点列表
        # 然后计算包围这些点的最小点集
        _, image_thresh = cv2.threshold(image_patch, self.THRESHOLD, 255, 1)
        image_thresh = cv2.Canny(image_thresh, 100, 100, 3)
        points = cv2.findNonZero(image_thresh)
        hull = cv2.convexHull(points)
    
        # 找到最小epsilon值，使近似后恰好得到4个点，范围通常在7到21之间
        # 这是包围图像的最小4点集
        for epsilon in range(3, 50):
            hull_simple = cv2.approxPolyDP(hull, epsilon, 1)
    
            if len(hull_simple) == 4:
                break
    
        hull = hull_simple
    
        # 找到最匹配的图像尺寸并进行透视变换/裁剪
        # （即最小化缩放比例）
    
        x,y,w,h = cv2.boundingRect(hull)
        target_corners = np.array([[0,0],[w,0],[w,h],[0,h]], np.float32)
    
        # 将凸包点排序为左上、右上、右下、左下顺序
        # 注：凸包已按顺时针排序，只需找到左上点的位置
    
        source_corners = hull.reshape(-1,2).astype('float32')
        min_dist = 100000
        index = 0
    
        for n in range(len(source_corners)):
            x,y = source_corners[n]
            dist = math.hypot(x,y)
    
            if dist < min_dist:
                index = n
                min_dist = dist
    
        # 旋转数组使左上点位于首位
        source_corners = np.roll(source_corners , -(2*index))
    
        try:
            transform = cv2.getPerspectiveTransform(source_corners, target_corners)
            return cv2.warpPerspective(image_patch, transform, (w,h))
    
        except:
            print ("变换失败", image_source)
            return image_patch

    def __repr__(self):
        return self.__class__.__name__+'()'

PyTorch神经网络搭建代码 #

# 构建ResNeXt模型:
transforms = v2.Compose([AdjustImage(),v2.Resize([256,256]), v2.PILToTensor()])
train_ds = datasets.ImageFolder("C:\\Users\\jenni\\Desktop\\Diss_Work\\X-ray_Images\\train", transform = transforms)
test_ds = datasets.ImageFolder("C:\\Users\\jenni\\Desktop\\Diss_Work\\X-ray_Images\\test", transform = transforms)
val_ds = datasets.ImageFolder("C:\\Users\\jenni\\Desktop\\Diss_Work\\X-ray_Images\\val", transform = transforms)

batch_size = 16
train_dataloader = DataLoader(train_ds, batch_size=batch_size)
test_dataloader = DataLoader(test_ds, batch_size=batch_size)
val_dataloader = DataLoader(val_ds, batch_size = batch_size)
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnext50_32x4d', pretrained=True)
print(model)
model.type(torch.LongTensor)
model.to(device)
loss_fn = nn.CrossEntropyLoss()
optimiser = torch.optim.Adam(model.parameters(), lr = 1e-3)

训练函数代码 #

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.type(torch.FloatTensor), y.type(torch.FloatTensor)
        X,y = X.to(device), y.to(device)
        # 计算预测误差
        pred = model(X)
        loss = loss_fn(pred, y.type(torch.LongTensor).to(device))

        # 反向传播
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

训练运行代码 #

epochs = 5
model.to(device)
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimiser)
    test(test_dataloader, model, loss_fn)
print("Done!")

尝试过的无效方案 #

• 单独转换图像时方法有效，但接入训练流程后问题复现
• 尝试使用np.asarray(image_source).astype(float)/255.0转换图像，保存后仍为全黑图像

最终解决办法 #

通过自定义ImageFolder的加载器解决PNG图像转换异常问题，加载器代码如下：

train_ds = datasets.ImageFolder(<path to folder>, transform = transforms, loader=lambda path: Image.fromarray((np.array(Image.open(path))>>8).astype(np.uint8)))

内容的提问来源于stack exchange，提问作者jennischofield02