machine-learning/captcha2/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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import os

import numpy
import cv2
import torch
from torch.utils.data import Dataset, DataLoader


char2id = {'0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8, '9': 9, 'a': 10, 'b': 11, 'c': 12,
           'd': 13, 'e': 14, 'f': 15, 'g': 16, 'h': 17, 'i': 18, 'j': 19, 'k': 20, 'l': 21, 'm': 22, 'n': 23, 'o': 24,
           'p': 25, 'q': 26, 'r': 27, 's': 28, 't': 29, 'u': 30, 'v': 31, 'w': 32, 'x': 33, 'y': 34, 'z': 35, 'A': 36,
           'B': 37, 'C': 38, 'D': 39, 'E': 40, 'F': 41, 'G': 42, 'H': 43, 'I': 44, 'J': 45, 'K': 46, 'L': 47, 'M': 48,
           'N': 49, 'O': 50, 'P': 51, 'Q': 52, 'R': 53, 'S': 54, 'T': 55, 'U': 56, 'V': 57, 'W': 58, 'X': 59, 'Y': 60,
           'Z': 61}

id2char = {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: 'a', 11: 'b', 12: 'c',
           13: 'd', 14: 'e', 15: 'f', 16: 'g', 17: 'h', 18: 'i', 19: 'j', 20: 'k', 21: 'l', 22: 'm', 23: 'n', 24: 'o',
           25: 'p', 26: 'q', 27: 'r', 28: 's', 29: 't', 30: 'u', 31: 'v', 32: 'w', 33: 'x', 34: 'y', 35: 'z', 36: 'A',
           37: 'B', 38: 'C', 39: 'D', 40: 'E', 41: 'F', 42: 'G', 43: 'H', 44: 'I', 45: 'J', 46: 'K', 47: 'L', 48: 'M',
           49: 'N', 50: 'O', 51: 'P', 52: 'Q', 53: 'R', 54: 'S', 55: 'T', 56: 'U', 57: 'V', 58: 'W', 59: 'X', 60: 'Y',
           61: 'Z'}


def label2tensor(label: str):
    t = torch.zeros(6)
    t[0] = char2id[label[0]]
    t[1] = char2id[label[1]]
    t[2] = char2id[label[2]]
    t[3] = char2id[label[3]]
    t[4] = char2id[label[4]]
    t[5] = char2id[label[5]]
    return t.long()


# def tensor2label(tensor: torch.Tensor):
#     lab = ''
#     lab += id2char[torch.argmax(tensor[:62]).item()]
#     lab += id2char[torch.argmax(tensor[62: 124]).item()]
#     lab += id2char[torch.argmax(tensor[124: 186]).item()]
#     lab += id2char[torch.argmax(tensor[186: 248]).item()]
#     lab += id2char[torch.argmax(tensor[248: 310]).item()]
#     lab += id2char[torch.argmax(tensor[310:]).item()]
#     return lab


def tensor2label(tensor: torch.Tensor):
    lab = ''
    lab += id2char[torch.argmax(tensor[0]).item()]
    lab += id2char[torch.argmax(tensor[1]).item()]
    lab += id2char[torch.argmax(tensor[2]).item()]
    lab += id2char[torch.argmax(tensor[3]).item()]
    lab += id2char[torch.argmax(tensor[4]).item()]
    lab += id2char[torch.argmax(tensor[5]).item()]
    return lab


def tensor2label2(tensor: torch.Tensor):
    lab = ''
    lab += id2char[tensor[0].item()]
    lab += id2char[tensor[1].item()]
    lab += id2char[tensor[2].item()]
    lab += id2char[tensor[3].item()]
    lab += id2char[tensor[4].item()]
    lab += id2char[tensor[5].item()]
    return lab


def transform(ndarray: numpy.ndarray):
    ndarray = ndarray / 255  # 不能写成 /=
    ndarray = numpy.transpose(ndarray, (2, 0, 1))
    return torch.tensor(ndarray, dtype=torch.float32)


class CaptchaDataset(Dataset):
    def __init__(self, kind='train'):
        self.path = os.path.join('dataset', kind)
        self.filenames = os.listdir(self.path)

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

    def __getitem__(self, item):
        image_name = self.filenames[item]
        image_path = os.path.join(self.path, image_name)
        image = cv2.imread(image_path)
        _X = transform(image)
        _y = image_name.split('.')[0]
        return _X, label2tensor(_y)


if __name__ == '__main__':
    dataset = CaptchaDataset(kind='validation')
    dataloader = DataLoader(dataset, batch_size=2)
    for X_batch, y_batch in dataloader:
        print(X_batch)
        print(y_batch)