from pvt_v2 import pvt_v2_b2
import cv2
import torchvision.transforms as transforms

# 读取测试图像
image = cv2.imread('E:/DeepLearning/PersonSearch/COAT/COAT/main/backbone/test.jpg')  # 替换为您的测试图像文件路径
image = cv2.resize(image, (224, 224))

# 使用 OpenCV 显示原始图像
cv2.imshow("Original Image", image)
cv2.waitKey(0)
cv2.destroyAllWindows()
# 创建转换以将图像缩放到 224x224 大小并转换为张量
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# 应用转换并将图像转换为张量
image_tensor = transform(image)
image_tensor = image_tensor.unsqueeze(0)  # 添加批次维度，将形状变为 [1, 3, 224, 224]

model = pvt_v2_b2(pretrained = "E:/DeepLearning/PersonSearch/COAT/COAT/main/backbone/pvt_v2_b2.pth")
# 使用模型进行推理
output = model(image_tensor)

print(output[0].shape)
print(output[1].shape)
print(output[2].shape)
print(output[3].shape)
# 在这里，您可以处理模型的输出，进行后续的操作或分析