machine-learning/captcha1/data.py at master · university-code-set/machine-learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import os

import cv2
import numpy
import torch
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split


def load_data():
    """
    X_train, X_test, y_train, y_test = load_data()
    """
    _X = range(5900)
    ndarray = numpy.loadtxt('./dataset/labels.csv', delimiter=',', dtype=str)
    _y = ndarray[:5900, -1]
    return train_test_split(_X, _y, train_size=0.7, test_size=0.3)


def label2array(label: str) -> torch.Tensor:
    """
    'ABCB' -> [1, 0, 0, ....] shape(26 * 4)
    """
    _y = torch.zeros(26 * 4)
    offset = 0
    _y[ord(label[0]) - 65 + offset] = 1
    offset += 26
    _y[ord(label[1]) - 65 + offset] = 1
    offset += 26
    _y[ord(label[2]) - 65 + offset] = 1
    offset += 26
    _y[ord(label[3]) - 65 + offset] = 1
    return _y


def array2label(array: torch.Tensor) -> str:
    """
    [1, 0, 0, ....] shape(26 * 4) -> 'ABCB'
    """
    label = ''
    label += chr(torch.argmax(array[:26]) + 65)
    label += chr(torch.argmax(array[26: 52]) + 65)
    label += chr(torch.argmax(array[52: 78]) + 65)
    label += chr(torch.argmax(array[78: 104]) + 65)
    return label


class CaptchaImageDataSet(Dataset):
    def __init__(self, X, y):  # noqa
        self.X = X
        self.y = y
        self.image_path = './dataset/images'

    def __len__(self):
        return len(self.X)

    def __getitem__(self, item):
        filename = os.path.join(self.image_path, str(self.X[item]) + '.jpg')
        ndarray = cv2.imread(filename) / 255  # 归一化
        ndarray = numpy.transpose(ndarray, (2, 0, 1))  # shape(3, 20, 80)
        _X = torch.tensor(ndarray, dtype=torch.float32)
        label = self.y[item]
        return _X, label2array(label)