YOLOv5 分类模型的预处理
flyfish
版本 6.2
将整个代码简化成如下代码
imgsz=224
file = "/home/a/Pictures/1.jpg"
transforms = classify_transforms(imgsz)
im = cv2.cvtColor(cv2.imread(file), cv2.COLOR_BGR2RGB)
print(im.shape)
im = transforms(im)
print(im.shape)
im = im.unsqueeze(0).to("cpu")
print(im.shape)
(511, 306, 3) H,W,C顺序
torch.Size([3, 224, 224]) 经过transforms后
torch.Size([1, 3, 224, 224]) 通过unsqueeze扩展增加一维,最后是NCHW的维度进入模型
图像经过了如下变换
def classify_transforms(size=224):
# Transforms to apply if albumentations not installed
return T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
最重要的是这两个
T.Resize(size)
T.CenterCrop(size)
分步演示
T.Resize(size) 图像的缩放
看一个参数还是两个参数
如果是一个参数int,那么图像的较小边将与该参数匹配,然后进行缩放,高宽比例不变。
如果是(h, w)
,那么图像就缩放到(h, w)
大小。
举个例子
如果 height > width 那么图片会被缩放到 (size * height / width, size).比例不变
参考:https://pytorch.org/vision/main/generated/torchvision.transforms.Resize.html
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import transforms
file_path = "./1.jpg"
img = Image.open(file)
print("Original:", img.size)
用了PIL库 是 宽w,高h 顺序
原始大小 宽w,高h Original: (306, 511)
trans0 = transforms.Compose([transforms.Resize(imgsz)])
after0 = trans0(img)
print("Resize:", after0.size)
after0.save('2.jpg')
经过Resize之后 宽w,高h Resize: (224, 374)
trans1 = transforms.Compose([transforms.CenterCrop(imgsz)])
after1 = trans1(after0)
print("CenterCrop:", after1.size)
after1.save('3.jpg')