视频帧（NumPy数组）转PyTorch张量后旋转出现颜色异常、黑屏问题求助

我太懂你这种折腾到崩溃的感觉了！ChatGPT越绕越乱的时候真的想摔键盘，咱们一步步拆解你代码里的问题，顺便解决你想直接旋转张量的需求～

先搞清楚两个核心问题的根源 #

1. 颜色异常：通道顺序+来回转换的逻辑混乱

你用cv2.VideoCapture读的帧是BGR格式，但PIL和PyTorch默认都是RGB格式，你转张量的时候没做通道转换，后续转PIL显示自然颜色全错。再加你的tensor_to_image函数里来回在PIL和cv2之间转，反而把通道绕得更乱了。

2. 黑屏：数值范围不匹配

你用transforms.ToTensor()把旋转后的PIL图转成张量时，这个函数会自动把0-255的像素值归一化到0-1的float类型，但你的check_tensor函数直接调用.byte()，等于把0-1的小数强制转成uint8，大部分值都会变成0，自然显示黑屏！

修正后的代码片段 #

第一步：修正帧转张量的逻辑（统一通道+数值范围）

读帧后转张量时，先把BGR转成RGB，同时统一数值范围（推荐用0-1的float，符合PyTorch的常用规范）：

def analyze_video(self, video_path):
    video_capture = cv2.VideoCapture(video_path)

    for i in range(1):
        ret, frame = video_capture.read()
        if not ret:
            break

        # 修正：把BGR转RGB，同时归一化到0-1的float
        frame_rgb = frame[...,::-1].copy()  # BGR转RGB
        frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0
        frame_tensor = frame_tensor.permute(2, 0, 1).unsqueeze(0)  # 转成(1,3,H,W)的张量格式

        self.check_tensor(frame_tensor)
  
        orig_prediction = self.analyze_frame(frame_tensor)
   
        rotated_tensor = im.rotate_tensor(frame_tensor, orig_prediction.head())
        self.check_tensor(rotated_tensor)

第二步：修正check_tensor函数（适配数值范围+通道转换）

显示时要把RGB转回BGR给cv2用，同时先把0-1的张量转回0-255的uint8：

def check_tensor(self, frame_tensor):
    # 先判断张量是0-1还是0-255范围，转回0-255的uint8
    if frame_tensor.dtype == torch.float32 and frame_tensor.max() <= 1.0:
        frame_numpy = (frame_tensor.squeeze(0).permute(1, 2, 0) * 255).byte().numpy()
    else:
        frame_numpy = frame_tensor.squeeze(0).permute(1, 2, 0).byte().numpy()
    
    # RGB转BGR适配cv2的显示要求
    frame_numpy = cv2.cvtColor(frame_numpy, cv2.COLOR_RGB2BGR)
    
    cv2.imshow("Frame", frame_numpy)
    cv2.waitKey(0)   
    cv2.destroyAllWindows()

第三步：修正rotate_tensor函数（解决颜色问题）

转PIL时先处理好张量的数值范围和通道，避免颜色错误：

def rotate_tensor(frame_tensor, landmarks):
    roll = calc_face_angle(landmarks)
    
    # 把0-1的张量转回RGB格式的uint8数组，再转PIL
    frame_np = (frame_tensor.squeeze(0).permute(1,2,0) * 255).byte().numpy()
    frame = Image.fromarray(frame_np)
    frame.show()  # 现在颜色应该正常了

    if not np.isnan(roll):
        rotated_frame = frame.rotate(roll, resample=Image.BICUBIC, expand=True)
    else:
        print("Failed to calculate face angle for rotation")
        return frame_tensor

    # 转回0-1的张量，保持和输入一致的格式
    transform = transforms.ToTensor()
    rotated_tensor = transform(rotated_frame).unsqueeze(0)
    return rotated_tensor

你的第二个问题：能不能直接旋转张量，不用转PIL？ #

当然可以！用torchvision.transforms.functional里的rotate函数直接操作张量，效率更高，还能避免来回转换的麻烦：

from torchvision.transforms.functional import rotate

def rotate_tensor_directly(frame_tensor, landmarks):
    roll = calc_face_angle(landmarks)
    if np.isnan(roll):
        print("Failed to calculate face angle for rotation")
        return frame_tensor
    
    # 直接旋转张量，expand=True会自动调整尺寸避免裁剪
    rotated_tensor = rotate(frame_tensor.squeeze(0), angle=roll, resample=Image.BICUBIC, expand=True).unsqueeze(0)
    return rotated_tensor

最后再划几个重点 #

• 时刻记住cv2是BGR，PyTorch/PIL是RGB，转换时一定要对应调整通道顺序
• 数值范围要统一：要么全程0-255的uint8，要么0-1的float，别混着来
• 直接操作张量比转PIL效率高，尤其是处理大量帧的时候

备注：内容来源于stack exchange，提问作者Мария Шипаева