yolov4数据预处理

最新推荐文章于 2023-04-05 13:48:08 发布

*Major*

最新推荐文章于 2023-04-05 13:48:08 发布

阅读量876

点赞数

本文链接：https://blog.csdn.net/qq_41375318/article/details/117336800

版权

import numpy as np
from PIL import Image
from torch.utils.data import Dataset

def get_data(annotation_line, input_shape):
    line = annotation_line.split()
    image = Image.open(line[0])
    iw, ih = image.size
    h, w = input_shape
    box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])


    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)
    dx = (w - nw) // 2
    dy = (h - nh) // 2

    image = image.resize((nw, nh), Image.BICUBIC)
    new_image = Image.new('RGB', (w, h), (128, 128, 128))
    new_image.paste(image, (dx, dy))
    image_data = np.array(new_image, np.float32)

    # 调整目标框坐标
    box_data = np.zeros((len(box), 5))
    if len(box) > 0:
        np.random.shuffle(box)
        box[:, [0, 2]] = box[:, [0, 2]] * nw / iw + dx
        box[:, [1, 3]] = box[:, [1, 3]] * nh / ih + dy
        box[:, 0:2][box[:, 0:2] < 0] = 0
        box[:, 2][box[:, 2] > w] = w
        box[:, 3][box[:, 3] > h] = h
        box_w = box[:, 2] - box[:, 0]
        box_h = box[:, 3] - box[:, 1]
        box = box[np.logical_and(box_w > 1, box_h > 1)]  # 保留有效框
        box_data = np.zeros((len(box), 5))
        box_data[:len(box)] = box

    return image_data, box_data

class LoadDataset(Dataset):
    # 确定数据路径
    def __init__(self, train_lines, image_size):
        super(LoadDataset, self).__init__()
        self.train_lines = train_lines
        self.train_batches = len(train_lines)
        self.image_size = image_size

    def __len__(self):
        return self.train_batches

    def __getitem__(self, index):
        lines = self.train_lines
        n = self.train_batches
        index = index % n

        img, y = get_data(annotation_line=lines[index], input_shape =self.image_size[0:2])
        # img = Image.fromarray(np.uint8(img))
        # img.show()
        if len(y) != 0:
            # 从坐标转换成0~1的百分比
            boxes = np.array(y[:, :4], dtype=np.float32)
            boxes[:, 0] = boxes[:, 0] / self.image_size[1]
            boxes[:, 1] = boxes[:, 1] / self.image_size[0]
            boxes[:, 2] = boxes[:, 2] / self.image_size[1]
            boxes[:, 3] = boxes[:, 3] / self.image_size[0]

            boxes = np.maximum(np.minimum(boxes, 1), 0)
            boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
            boxes[:, 3] = boxes[:, 3] - boxes[:, 1]

            boxes[:, 0] = boxes[:, 0] + boxes[:, 2] / 2
            boxes[:, 1] = boxes[:, 1] + boxes[:, 3] / 2
            y = np.concatenate([boxes, y[:, -1:]], axis=-1)

        img = np.array(img, dtype=np.float32)

        tmp_inp = np.transpose(img / 255.0, (2, 0, 1))
        tmp_targets = np.array(y, dtype=np.float32)
        return tmp_inp, tmp_targets


if __name__ == '__main__':
    import torchvision.transforms as transforms
    from torch.utils.data import DataLoader

    norm_mean = [0.33424968, 0.33424437, 0.33428448]
    norm_std = [0.24796878, 0.24796101, 0.24801227]

    train_transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),  # 0-255 归一化到0-1 转Tensor
        transforms.Normalize(norm_mean, norm_std),
    ])
    annotation_path = r'D:\paddy_detect\yolov4-pytorch-master\2009_train.txt'

    with open(annotation_path) as f:
        lines = f.readlines()

    train_dataset = LoadDataset(train_lines=lines[:10], image_size=(608, 608))
    train_loader = DataLoader(dataset=train_dataset, batch_size=1, shuffle=True)  # shuffle训练时打乱样本

    for idx, data_info in enumerate(train_loader):
        print(idx)
        print(data_info)

‘D:\paddy_detect\yolov4-pytorch-master\2009_train.txt’

D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/0.jpg 11,45,51,117,0 62,52,114,122,0 108,1,153,31,0 155,195,229,227,0 2,194,27,267,0 219,70,249,142,0 176,32,216,117,0 193,18,265,51,0 15,270,44,299,0 35,208,102,225,0 217,156,290,230,0 280,153,299,197,0 55,265,130,283,0 141,57,174,87,0 123,78,176,108,0 281,1,300,29,0 245,51,264,88,0 2,159,46,186,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/1.jpg 60,45,104,107,0 88,10,180,43,0 110,142,151,213,0 141,167,201,205,0 1,170,32,201,0 1,213,31,279,0 162,118,198,157,0 173,99,212,118,0 119,68,184,102,0 203,1,229,50,0 236,29,254,102,0 230,1,258,24,0 268,90,299,134,0 220,126,260,162,0 252,251,286,265,0 271,259,299,298,0 183,244,235,299,0 152,223,225,252,0 76,263,171,292,0 79,225,132,261,0 95,222,175,266,0 1,50,26,77,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/10.jpg 233,178,275,233,0 207,150,235,194,0 149,124,240,144,0 284,72,299,143,0 250,76,289,115,0 212,77,253,103,0 212,1,244,39,0 235,31,263,49,0 272,238,299,252,0 269,252,298,297,0 72,173,128,210,0 96,215,152,263,0 1,206,57,228,0 1,231,30,251,0 2,254,37,299,0 102,267,120,300,0 59,284,98,298,0 37,150,83,178,0 27,157,73,202,0 175,1,196,70,0 106,1,142,29,0 49,1,86,9,0 29,1,71,41,0 10,13,45,43,0 50,95,86,141,0 78,112,123,136,0 64,88,91,114,0 103,71,141,110,0 149,83,201,118,0 281,174,300,196,0 281,199,299,217,0 102,51,172,80,0 102,25,144,57,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/11.jpg 78,195,162,227,0 32,131,58,173,0 8,88,40,120,0 165,102,236,135,0 233,144,265,162,0 247,101,274,146,0 272,133,299,205,0 247,191,289,212,0 219,215,295,278,0 232,4,261,37,0 268,7,300,30,0 1,1,62,28,0 1,45,17,74,0 19,232,69,259,0 15,254,68,297,0 131,247,175,288,0 167,214,210,278,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/12.jpg 17,1,44,39,0 1,1,15,45,0 1,55,46,93,0 2,133,40,180,0 15,176,99,195,0 85,111,130,127,0 109,97,160,171,0 104,184,154,220,0 41,241,65,298,0 130,169,206,266,0 94,262,137,292,0 96,228,132,254,0 245,146,276,201,0 210,145,245,183,0 264,81,300,136,0 254,47,299,90,0 236,41,254,108,0 203,62,235,111,0 151,85,205,141,0 144,275,212,299,0 155,1,296,17,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/13.jpg 1,84,62,117,0 71,121,108,186,0 51,130,81,152,0 19,165,81,221,0 1,120,22,140,0 134,245,168,298,0 196,226,222,275,0 219,157,242,224,0 256,22,270,112,0 169,55,189,134,0 214,47,240,114,0 148,1,168,78,0 288,2,298,68,0 91,262,114,299,0 43,231,83,291,0 5,205,50,254,0 2,251,41,293,0 259,160,290,210,0 235,233,258,288,0 273,272,300,290,0 53,3,107,52,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/14.jpg 100,61,136,148,0 1,173,51,210,0 27,236,59,283,0 113,223,137,254,0 150,262,176,295,0 187,182,233,252,0 255,147,289,203,0 216,204,269,264,0 141,136,165,195,0 158,191,184,233,0 1,13,57,51,0 24,73,63,106,0 6,115,52,144,0 68,95,83,149,0 181,123,224,162,0 240,122,270,153,0 270,100,300,113,0 239,88,268,103,0 190,3,210,66,0 94,1,115,43,0 249,268,281,300,0 242,27,294,78,0 144,2,159,44,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/15.jpg 78,69,128,136,0 170,59,211,101,0 243,137,264,190,0 125,201,181,253,0 1,196,62,229,0 80,250,150,273,0 55,274,141,298,0 15,221,31,293,0 29,221,50,278,0 283,195,299,258,0 288,58,299,89,0 244,26,295,50,0 204,15,223,59,0 86,1,170,29,0 37,31,57,67,0 1,37,35,78,0 235,1,299,22,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/16.jpg 241,1,260,30,0 286,5,299,59,0 190,1,206,25,0 227,119,244,151,0 241,238,261,284,0 22,112,98,154,0 1,87,43,158,0 2,31,29,96,0 58,22,94,36,0 44,7,59,46,0 25,25,55,48,0 101,63,160,81,0 85,85,181,107,0 223,41,242,68,0 175,70,239,99,0 177,34,224,80,0 265,125,296,157,0 106,132,163,178,0 119,259,158,298,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/17.jpg 91,148,136,175,0 83,167,121,203,0 161,199,188,257,0 195,78,268,124,0 247,141,300,181,0 252,209,299,234,0 75,43,129,112,0 6,105,33,161,0 19,186,84,220,0 39,268,104,288,0 241,2,283,63,0 176,32,191,62,0 112,13,170,66,0 48,5,101,43,0 31,26,75,76,0 5,27,32,74,0 3,2,25,32,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/18.jpg 2,8,38,90,0 107,114,195,144,0 201,86,254,155,0 140,169,211,190,0 130,225,156,285,0 96,191,124,248,0 39,187,93,259,0 31,170,76,224,0 1,131,55,197,0 43,116,77,167,0 58,257,84,297,0 221,140,238,229,0 245,173,296,221,0 268,130,299,173,0 264,211,298,247,0 174,235,189,298,0 213,224,281,295,0 107,5,140,103,0 261,24,298,67,0 264,60,289,88,0 69,41,89,99,0 176,191,205,219,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/19.jpg 138,42,167,170,0 123,248,147,291,0 63,255,112,298,0 42,194,90,251,0 1,142,18,195,0 235,209,269,300,0 191,189,251,221,0 186,148,239,187,0 209,129,238,152,0 255,132,278,203,0 269,164,299,239,0 54,43,134,83,0 180,13,259,72,0 257,35,297,86,0 1,231,35,295,0 32,2,82,18,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/2.jpg 173,1,241,26,0 51,60,134,118,0 15,31,96,51,0 2,38,25,62,0 265,31,286,96,0 250,93,285,136,0 143,103,215,122,0 166,125,211,175,0 207,42,266,64,0 216,231,274,294,0 193,255,218,294,0 92,240,150,280,0 160,281,182,298,0 69,267,105,299,0 80,149,127,169,0 142,82,189,103,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/3.jpg 55,117,95,153,0 1,84,20,133,0 32,184,62,204,0 11,139,57,170,0 91,212,163,240,0 160,200,215,279,0 212,118,271,188,0 185,101,273,130,0 248,68,298,114,0 234,4,267,55,0 126,13,231,28,0 32,3,81,37,0 14,1,31,30,0 12,276,45,299,0 66,282,93,298,0 238,245,270,286,0 283,204,299,246,0 244,176,281,220,0 197,218,266,241,0 29,62,113,85,0 63,46,115,71,0 108,51,175,86,0 112,91,183,115,0 83,239,132,266,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/4.jpg 48,82,70,134,0 130,37,221,104,0 259,121,298,139,0 278,147,299,161,0 132,152,167,177,0 122,170,213,219,0 238,227,271,276,0 112,229,159,271,0 41,230,115,291,0 1,204,27,262,0 1,67,17,84,0 174,264,233,298,0 279,246,298,296,0 215,195,232,231,0 169,138,231,172,0 217,173,266,200,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/5.jpg 186,157,210,175,0 250,165,271,210,0 242,211,261,234,0 262,230,281,253,0 286,56,300,129,0 195,90,227,150,0 137,77,205,125,0 102,75,140,122,0 86,22,160,52,0 132,2,167,37,0 98,1,135,24,0 240,1,257,67,0 210,57,227,87,0 220,59,249,89,0 65,153,107,190,0 228,250,265,298,0 143,255,220,274,0 155,204,204,255,0 124,241,167,276,0 64,201,147,222,0 90,238,116,298,0 17,246,38,299,0 31,183,55,252,0 44,243,84,271,0 28,43,58,145,0 2,139,28,161,0 1,121,25,137,0 262,131,277,150,0 23,1,72,35,0 2,1,29,26,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/6.jpg 25,114,48,183,0 47,180,77,221,0 75,216,100,256,0 111,196,173,271,0 208,238,261,292,0 245,259,265,298,0 180,243,232,279,0 167,201,228,220,0 132,148,197,179,0 162,137,193,163,0 179,129,231,173,0 275,86,297,167,0 223,104,272,151,0 189,71,239,94,0 135,4,156,59,0 269,3,298,46,0 214,2,246,44,0 244,6,269,35,0 184,10,224,35,0 90,104,157,136,0 77,83,170,94,0 61,20,133,46,0 60,1,81,25,0 61,61,98,82,0 49,27,123,68,0 25,262,78,299,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/7.jpg 53,71,122,131,0 119,90,148,175,0 100,190,128,258,0 35,255,111,270,0 225,38,300,52,0 214,52,294,99,0 259,143,299,175,0 254,184,300,210,0 262,165,291,244,0 2,59,48,86,0 1,114,26,141,0 105,1,131,36,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/8.jpg 84,145,115,219,0 126,191,192,259,0 89,197,149,246,0 127,125,165,166,0 205,104,241,130,0 201,144,248,173,0 269,105,298,120,0 252,276,294,298,0 39,269,83,298,0 1,267,38,292,0 44,220,65,262,0 66,195,81,275,0 19,113,54,143,0 42,127,67,146,0 33,51,94,75,0 5,63,67,83,0 1,73,30,113,0 12,20,81,42,0 79,1,111,52,0 1,6,23,62,0 206,1,253,44,0 275,13,296,38,0 279,39,299,85,0
D:\paddy_detect\yolov4-pytorch-master/VOCdevkit/VOC2009/JPEGImages/9.jpg 1,142,49,167,0 60,170,95,216,0 99,262,128,300,0 159,224,172,298,0 203,215,223,278,0 177,262,207,299,0 180,195,199,261,0 149,108,167,178,0 184,138,248,167,0 149,35,198,57,0 105,1,136,39,0 67,72,105,131,0 102,83,151,104,0 1,74,51,95,0 36,95,79,145,0 1,42,18,72,0 177,62,234,143,0 233,64,296,91,0 173,1,239,24,0 197,27,231,47,0 276,115,288,149,0 286,130,299,156,0