[Python] Problem with CNN model for prediction of color from an black and white image

I've created a CNN model, which gets as an input 13x13 images grayscaled to LAB L-channel. I trained it with 12000 images and got an MSE error of 7. To use the model, I need to slice an image into 13x13 little images and feed them through the network. From that of I take the L-channel from the orginal little image and put it with the two predictions (A- and B-channel) together to one pixel for the predicted image with colors. After that I put all these predicted pixels together and form the predicted image. This is the function, that creates these predicted pixels:

def create_pxl_from_preds(input_image, prediction):
prediction_rescaled = torch.mul(prediction, 128)
a, b = prediction_rescaled[0]
a = a.detach().numpy()
b = b.detach().numpy()

l_channel = input_image[6, 6]

a_channel = np.full((1, 1), a, dtype=np.float32)
b_channel = np.full((1, 1), b, dtype=np.float32)
l_channel = np.full((1, 1), l_channel, dtype=np.float32)

l_channel = (l_channel / 255 * 100).astype(np.float32)
a_channel = a_channel.astype(np.float32)
b_channel = b_channel.astype(np.float32)

image_pred = cv2.merge([l_channel, a_channel, b_channel])
return image_pred


After that I have a function to build the rows of the image and a function to put these rows together and form a whole image. Also there are always 487 images for one row and also together 487 rows

def rebuild_image_pxl_row(
start_calc: int,
end_calc: int,
num_sections_per_row=487,
) -> None:
tensor_rows = torch.arange(start_calc, end_calc)
list_rows = tensor_rows.tolist()
add_index = 0
for i in list_rows:
idx = i

image_files = [f for f in os.listdir(temp_folder_images) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
image_files.sort(key=extract_numbers)

start_idx = idx * num_sections_per_row
start_idx += add_index

end_idx = start_idx + num_sections_per_row
row_sections = image_files[start_idx:end_idx]
print(f"start idx: {start_idx}, end idx: {end_idx}")

row_images = []

for index, section_file in enumerate(row_sections):
section_path = os.path.join(temp_folder_images, section_file)
section_image = cv2.imread(section_path, cv2.IMREAD_GRAYSCALE)
section_tensor = torch.from_numpy(section_image).float().unsqueeze(0).unsqueeze(0)
section_pred = conv_model(section_tensor)
section_reconstructed = create_pxl_from_preds(section_image, section_pred)
row_images.append(np.round(cv2.cvtColor(section_reconstructed, cv2.COLOR_LAB2BGR)*255.0, 0))

row = np.hstack(row_images)
row_path = os.path.join(temp_folder_rows, f"row_{idx}.png")
cv2.imwrite(row_path, row)
add_index += 1
print(f"Row {idx} reconstructed and saved as {row_path}.")


def rebuild_image_pxl(row_ordner, target_height=487):
row_files = [f for f in os.listdir(row_ordner) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
row_files.sort(key=extract_numbers)
print(row_files[:20])

row_files = row_files[:target_height]

all_rows = []

for index, row_file in enumerate(row_files):
row_path = os.path.join(row_ordner, row_file)
row_image = cv2.imread(row_path)
all_rows.append(row_image)

final_image = cv2.vconcat(all_rows)
final_image = np.fliplr(final_image)
final_image = cv2.rotate(final_image, cv2.ROTATE_90_COUNTERCLOCKWISE)
cv2.imwrite("image.png", final_image)
show_image("image.png")


I use opencv as a image processing library. But when I try to run the program no errors occur, but the colors of the predicted image is not right and not even close to the one of the original one. The edges and nuances of the L-Channel are detected, so the copying of this data works. But the color prediction doesn't and I think it's something with opencv and the different color spaces BGR, RGB and LAB. Because opencv opens and saves images in BGR color space, I just converted the image values from LAB to BGR. I need help figuring out, where the error is.

The predicted image looks like this: enter image description here

Continue reading...