Почему мой тензор изображения выводит все нули после обработки с помощью Spatial Transformer Network Pytorch? ⇐ Python

[Anonymous]

Я работаю над небольшим проектом, включающим сеть пространственных преобразователей (STN) для обработки изображений. Я случайно загрузил ветку с непроверенным кодом, и теперь столкнулся с проблемой, из-за которой выходные данные тензора изображения — все нули.
Вот соответствующая часть моего кода:

Код: Выделить всё

import torch
from torchvision import transforms
from PIL import Image
from pathlib import Path
from model import STModel
from typing import Union
import numpy as np

class STN:
"""
Class to handle the processing of a single image using a Spatial Transformer Network (STN).

Args:
pretrained (Path): Path to the pre-trained model.
"""

def __init__(self, pretrained: Union[str, Path]) -> None:
self.pretrained: Path = pretrained
self.device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model: STModel = STModel().to(self.device)
self.model.load_state_dict(torch.load(self.pretrained, map_location=self.device))
self.model.eval()

self.transform: transforms.Compose = transforms.Compose([
transforms.Resize((150, 120)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])  # Adjust normalization for single-channel input
])

def process_image(self, input_path: Union[str, Path], output_path: Union[str, Path]) -> None:
"""
Process a single image using the pre-trained model.

Args:
input_path (Union[str, Path]): Path to the input image.
output_path (Union[str, Path]): Path where the output image will be saved.
"""
input_path: Path = Path(input_path)
output_path: Path = Path(output_path)
image: Image.Image = Image.open(input_path).convert('L')  # Ensure the image is in greyscale
print(f"Loaded image: {input_path}")
print(f"Image size: {image.size}")
print(f"Image mode: {image.mode}")

input_tensor: torch.Tensor = self.transform(image).unsqueeze(0).to(self.device)
print(f"Transformed tensor shape: {input_tensor.shape}")
print(f"Transformed tensor min, max: {input_tensor.min().item()}, {input_tensor.max().item()}")

with torch.no_grad():
output_tensor: torch.Tensor = self.model(input_tensor)
print(f"Output tensor shape: {output_tensor.shape}")
print(f"Output tensor min, max: {output_tensor.min().item()}, {output_tensor.max().item()}")

output_array = np.array([output_tensor.squeeze().cpu().detach()])

print(f"Processed and saved output image: {output_path}")
print(f"Output image content: {output_array}")
print(f"Output tensor shape: {output_tensor.shape}")

if __name__ == "__main__":
stn: STN = STN(pretrained="spt_model.pt")
stn.process_image(
input_path=Path("dataset/train/aaAGoBxqnJgoEGzD.jpg"),
output_path=Path("output.jpg")
)

Когда я запускаю код, выходной_тензор имеет минимальное и максимальное значения (0,0, 0,0), и когда я распечатываю выходной_массив , там естественно тоже все нули. Вот образец распечатки:

Код: Выделить всё

Loaded image: dataset/train/aaAGoBxqnJgoEGzD.jpg
Image size: (150, 120)
Image mode: L
Transformed tensor shape: torch.Size([1, 1, 150, 120])
Transformed tensor min, max: -2.8582..., 1.5927...
Output tensor shape: torch.Size([1, 1, 150, 120])
Output tensor min, max: 0.0, 0.0
Processed and saved output image: output.jpg
Output image content: [[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
...
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]
[0. 0. 0. ... 0. 0. 0.]]
Output tensor shape: torch.Size([1, 1, 150, 120])

Я подозреваю, что может возникнуть проблема с тем, как изображение преобразуется или как модель обрабатывает изображение. Буду очень признателен за любую информацию о том, почему это может происходить и как это исправить.
Вы можете найти полный код проекта на моем GitHub для получения дополнительной информации: Номерной знак STN.

Подробнее здесь: https://stackoverflow.com/questions/786 ... transforme