Google Colab中如何让cv2_imshow输出图像的右侧显示print打印内容

实现思路 #

Google Colab支持直接渲染HTML内容，我们可以通过flex布局将图像、文本放在同一行的两个容器中，不需要调整原有模型预测逻辑，仅需修改输出部分的代码即可实现需求。

前置依赖导入 #

import cv2
import base64
import numpy as np
from IPython.display import display, HTML

修改后完整函数代码 #

def ensemble_predict(scheme_model,quality_model,rotate_model,img,y=None):
    quality_labels  = ["blur","sharp","torch"]
    rotate_labels = [0,90,-90,180]
    bgr_labels = ["bgr","rgb"]
  
    scheme_pred = scheme_model.predict(img)[0]
    quality_pred = quality_model.predict(img)[0]
    rotate_pred = rotate_model.predict(img)[0]
    # 处理待显示的图像
    show_img = cv2.resize(img[0]*255,(256,256)).astype(np.uint8)
    
    # 先拼接所有待输出的文本
    p = scheme_pred
    scheme_pred = 0 if scheme_pred < 0.5 else 1 
    text_content = ""
    text_content += f"Predict BGR or RGB: {bgr_labels[scheme_pred]}, p: {p[0]}\n"
    text_content += f"Predict quality: {quality_labels[np.argmax(quality_pred)]}, p: {quality_pred[np.argmax(quality_pred)]}\n"
    text_content += f"Predict rotation: {rotate_labels[np.argmax(rotate_pred)]}, p: {rotate_pred[np.argmax(rotate_pred)]}\n\n"
    
    if y is not None:
        q_l = quality_labels[y[0]]
        r_l = rotate_labels[y[1]]
        b_l = bgr_labels[y[2]]
        text_content += f"True BGR or RGB: {b_l}\n"
        text_content += f"True quality: {q_l}\n"
        text_content += f"True rotation: {r_l}\n"

    # 将cv2图像转为base64编码
    _, encoded_img = cv2.imencode('.png', show_img)
    base64_img = base64.b64encode(encoded_img).decode('utf-8')

    # 构建并排布局的HTML
    html_str = f"""
    <div style="display:flex;align-items:center;gap:20px;">
        <img src="data:image/png;base64,{base64_img}" style="width:256px;height:256px;">
        <pre style="margin:0;font-size:14px;line-height:1.5;">{text_content}</pre>
    </div>
    """
    # 渲染输出
    display(HTML(html_str))

替代方案（Matplotlib实现） #

如果不想用HTML渲染，也可以用Matplotlib的子图实现并排效果，代码示例如下：

import matplotlib.pyplot as plt

def ensemble_predict_plt(scheme_model,quality_model,rotate_model,img,y=None):
    quality_labels  = ["blur","sharp","torch"]
    rotate_labels = [0,90,-90,180]
    bgr_labels = ["bgr","rgb"]
  
    scheme_pred = scheme_model.predict(img)[0]
    quality_pred = quality_model.predict(img)[0]
    rotate_pred = rotate_model.predict(img)[0]
    # matplot默认显示RGB顺序，需要转换cv2的BGR格式
    show_img = cv2.cvtColor(cv2.resize(img[0]*255,(256,256)).astype(np.uint8), cv2.COLOR_BGR2RGB)
    
    # 拼接文本
    p = scheme_pred
    scheme_pred = 0 if scheme_pred < 0.5 else 1 
    text_content = ""
    text_content += f"Predict BGR or RGB: {bgr_labels[scheme_pred]}, p: {p[0]}\n"
    text_content += f"Predict quality: {quality_labels[np.argmax(quality_pred)]}, p: {quality_pred[np.argmax(quality_pred)]}\n"
    text_content += f"Predict rotation: {rotate_labels[np.argmax(rotate_pred)]}, p: {rotate_pred[np.argmax(rotate_pred)]}\n\n"
    
    if y is not None:
        q_l = quality_labels[y[0]]
        r_l = rotate_labels[y[1]]
        b_l = bgr_labels[y[2]]
        text_content += f"True BGR or RGB: {b_l}\n"
        text_content += f"True quality: {q_l}\n"
        text_content += f"True rotation: {r_l}\n"

    # 绘制并排布局
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
    ax1.imshow(show_img)
    ax1.axis('off')
    ax2.text(0, 0.5, text_content, fontsize=12, va='center', linespacing=1.5)
    ax2.axis('off')
    plt.tight_layout()
    plt.show()

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