def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, '*.png'))
targets_orig = [x.split('/')[-1][:-4] for x in image_files]
targets = [[c for c in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
label_enc = preprocessing.LabelEncoder()
label_enc.fit(targets_flat)
targets_enc = [label_enc.transform(x) for x in targets]
targets_enc = np.array(targets_enc) + 1
train_imgs, test_imgs, train_targets, test_targets, train_orig_targets, test_orig_targets = model_selection.train_test_split(
image_files, targets_enc, targets_orig, test_size=0.1, random_state=42)
train_dataset = dataset.Classification(image_paths=train_imgs,
targets=train_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.BATCH_SIZE,
shuffle=True)
test_dataset = dataset.Classification(image_paths=test_imgs,
targets=test_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=config.BATCH_SIZE,
shuffle=False)
model = CaptchaModel(num_chars=len(label_enc.classes_))
model = model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=5,
verbose=True)
for epoch in range(config.EPOCHS):
train_loss = engine.train(model, train_loader, optimizer)
val_preds, valid_loss = engine.eval(model, test_loader)
print(
f"Epoch: {epoch}: Train loss: {train_loss}, Valid loss: {valid_loss}"
)
valid_cap_preds = []
for vp in val_preds:
current_preds = decode_predictions(vp, label_enc)
valid_cap_preds.extend(current_preds)
print(list(zip(test_orig_targets, valid_cap_preds))[6:11])
def evaluate():
files = []
for img_ext in config.ALLOWED_EXTENSIONS:
files.extend(
glob(
os.path.join(config.IMAGES_UPLOADED_PATH,
"*.{}".format(img_ext))))
files.sort(key=os.path.getctime, reverse=True)
test_img = files[:1]
test_dataset = dataset.ClassificationDataset(image_paths=test_img,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
test_loader = DataLoader(
test_dataset,
batch_size=1,
num_workers=0,
)
model = CaptchaModel(num_chars=len(lbl_enc.classes_))
model.to(config.DEVICE)
model.load_state_dict(
load('./checkpoints/captcha_v1/captcha_v1.pth',
map_location=config.DEVICE))
model.eval()
for data in test_loader:
data["images"] = data["images"].to(config.DEVICE)
prediction, _ = model(data["images"])
prediction_output = decode_predictions(prediction, lbl_enc)
return prediction_output
def predict(captcha,
model_dir='./model/model-latest.pkl',
use_gpu=True,
mode='captcha'):
"""
:param captcha:
:param model_dir:
:param use_gpu:
:param mode:
:return:
"""
gpu_available = torch.cuda.is_available()
if mode == 'captcha':
from model import CaptchaModel
elif mode == 'kaptcha':
from kaptcha_model import CaptchaModel
else:
return
model = CaptchaModel()
if use_gpu and gpu_available:
model_state = torch.load(model_dir)
else:
model_state = torch.load(model_dir,
map_location=lambda storage, loc: storage)
model.load_state_dict(model_state['network'])
if use_gpu and gpu_available:
model = model.cuda()
else:
model = model.cpu()
transformer = Compose(ToTensor())
img_pil = Image.open(captcha)
img_tensor = transformer.transforms(img_pil)
model.eval()
x = torch.stack([img_tensor])
if use_gpu and gpu_available:
x = x.cuda()
pred1, pred2, pred3, pred4 = model(x)
pred_seq = [
torch.argmax(pred1).item(),
torch.argmax(pred2).item(),
torch.argmax(pred3).item(),
torch.argmax(pred4).item()
]
pred_seq = [item + 1 for item in pred_seq]
_, id2label = get_dict()
res = ''.join([id2label[i] for i in pred_seq])
return res
def run_training():
image_files = glob.glob(
os.path.abspath(os.path.join(config.DATA_DIR, "*.png")))
labels = [list(x.split("/")[-1].split(".")[0]) for x in image_files]
labels_flat = [c for x in labels for c in x]
label_enc = preprocessing.LabelEncoder()
label_enc.fit(labels_flat)
tar_enc = np.array([label_enc.transform(x) for x in labels]) + 1
train_X, test_X, train_y, test_y, train_target, test_target = model_selection.train_test_split(
image_files, tar_enc, labels)
train_dataset = dataset.DataSet(train_X,
train_y,
resize=(config.IMG_HEIGHT,
config.IMG_WIDTH))
test_dataset = dataset.DataSet(test_X,
test_y,
resize=(config.IMG_HEIGHT,
config.IMG_WIDTH))
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True)
test_dataloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS)
cm = CaptchaModel(num_chars=len(label_enc.classes_))
cm.to(config.DEVICE)
optimizer = torch.optim.Adam(cm.parameters(), lr=3e-5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=10,
verbose=True)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(cm, train_dataloader, optimizer)
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR,"*.png"))
# path to the dataset
targets_orig = [x.split('/')[-1][:-4] for x in image_files]
targets = [[c for c in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
lbl_enc = preprocessing.LabelEncoder()
lbl_enc.fit(targets_flat)
targets_enc = [lbl_enc.transform(x) for x in targets]
targets_enc = np.array(targets_enc) + 1
# print(targets)
# print(target_enc)
# print(len(lbl_enc.classes_))
# # print(targets_orig)
# for i, item in enumerate(lbl_enc.classes_):
# print(item, '-->', i)
train_imgs, test_imgs, train_targets, test_targets, train_orig_targets, test_orig_targets= model_selection.train_test_split(image_files, targets_enc, targets_orig, test_size = 0.1, random_state= 42)
train_dataset = dataset.ClassificationDataset(image_paths = train_imgs, targets = train_targets, resize = (config.IMAGE_HEIGHT, config.IMAGE_WIDTH))
print(train_dataset[0])
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size = config.BATCH_SIZE,
num_workers = config.NUM_WORKERS,
shuffle = True
)
test_dataset = dataset.ClassificationDataset(image_paths = test_imgs, targets = test_targets, resize = (config.IMAGE_HEIGHT, config.IMAGE_WIDTH))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size = config.BATCH_SIZE,
num_workers = config.NUM_WORKERS,
shuffle = False
)
model = CaptchaModel(num_chars = len(lbl_enc.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr = 3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor =0.8, patience= 5, verbose= True)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
valid_pred, valid_loss = engine.eval_fn(model, train_loader)
def load():
img_files = glob.glob(os.path.join(config.data_dir, "*.png"))
labels_orig = [x.split('/')[-1][:-4] for x in img_files]
labels = [[char for char in x] for x in labels_orig] # all len 5
lab_flat = [char for clist in labels for char in clist]
encoder = LabelEncoder()
encoder.fit_transform(lab_flat)
targets_enc = [encoder.transform(x) for x in labels]
targets_enc = np.array(targets_enc) + 1
train_imgs, test_imgs, train_lab, test_lab, train_orig_lab, test_orig_lab = train_test_split(
img_files, targets_enc, labels_orig, test_size=0.1)
train_ds = dataset.Classification(train_imgs,
targets=train_lab,
resize=(config.h, config.w))
train_loader = DataLoader(train_ds,
shuffle=True,
batch_size=config.batch_size,
num_workers=config.workers)
test_ds = dataset.Classification(test_imgs,
targets=test_lab,
resize=(config.h, config.w))
test_loader = DataLoader(test_ds,
shuffle=False,
batch_size=config.batch_size,
num_workers=config.workers)
model = CaptchaModel(len(encoder.classes_)).to(config.device)
optimizer = Adam(model.parameters())
sched = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=5,
verbose=True)
for epoch in range(10):
train_loss = utils.train(model, train_loader, optimizer)
valid_pred, valid_loss = utils.eval(model, train_loader)
print(
f"Epoch: {epoch}, train_loss: {train_loss}, val_loss: {valid_loss}"
)
def get_predictions(image_path, model_path):
classes = [
'2', '3', '4', '5', '6', '7', '8', 'b', 'c', 'd', 'e', 'f', 'g', 'm',
'n', 'p', 'w', 'x', 'y'
]
le = preprocessing.LabelEncoder()
le.fit(sorted(classes))
n_classes = len(classes)
model = CaptchaModel(num_chars=n_classes)
model.load_state_dict(torch.load(model_path))
model.eval()
data = preproc_image(image_path)
with torch.no_grad():
preds, _ = model(**data)
# Now decode the preds
preds = decode_predictions(preds, le)
preds = remove_blanks(preds)
print(preds)
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
targets_orig = [x.split("\\")[-1][:-4] for x in image_files]
targets = [[y for y in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
label_enc = preprocessing.LabelEncoder()
label_enc.fit(targets_flat)
targets_enc = [label_enc.transform(x) for x in targets]
targets_enc = np.array(targets_enc) + 1
# print(targets_enc)
# print(label_enc.classes_)
(
train_imgs,
test_imgs,
train_targets,
test_targets,
train_orig_targets,
test_orig_targets,
) = model_selection.train_test_split(image_files,
targets_enc,
targets_orig,
test_size=0.1,
random_state=42)
train_dataset = dataset.ClassificationDataset(
image_paths=train_imgs,
targets=train_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH),
)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True,
pin_memory=True)
test_dataset = dataset.ClassificationDataset(
image_paths=test_imgs,
targets=test_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH),
)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False,
pin_memory=True)
model = CaptchaModel(num_chars=len(label_enc.classes_)).cuda()
model.to(torch.device(config.DEVICE))
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=5,
verbose=True)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
valid_preds, valid_loss = engine.eval_fn(model, test_loader)
valid_cap_preds = []
for vp in valid_preds:
current_preds = decode_predictions(vp, label_enc)
valid_cap_preds.extend(current_preds)
pprint.pprint(list(zip(test_orig_targets, valid_cap_preds))[:10])
test_dup_rem = [remove_duplicates(c) for c in test_orig_targets]
accuracy = metrics.accuracy_score(test_dup_rem, valid_cap_preds)
print(
f"EPOCH: {epoch}.train_loss:{train_loss},valid_loss:{valid_loss}, Accuracy={accuracy}"
)
scheduler.step(valid_loss)
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
# "/../../azopekr.png"
targets_orig = [x.split("/")[-1][:4] for x in image_files]
# abcde -> [a, b, c, d, e]
"""
targets
['6', 'd', 'm', 'x'],
['c', '7', '5', '3'],
['g', 'g', 'd', '7'],
['x', 'e', 'm', 'y'],
['6', 'g', '4', '5'],
['p', '2', 'x', '7'],
['d', 'y', 'p', '7'],
['6', 'e', 'c', 'b'],
['3', 'm', 'x', 'd'],
['f', 'c', 'm', 'e'],
['8', 'n', '6', '2'],
"""
targets = [[c for c in x] for x in targets_orig]
"""
targets_flat
['e',
'2',
'd',
'6',
'f',
'w',
'3',
'b',
'n',
...
]
"""
targets_flat = [c for clist in targets for c in clist]
lbl_enc = preprocessing.LabelEncoder()
lbl_enc.fit(targets_flat)
targets_enc = [lbl_enc.transform(x) for x in targets]
# I added one because 0 is kept for unknown
targets_enc = np.array(targets_enc) + 1
print(targets_enc)
print(len(lbl_enc.classes_))
train_imgs, test_imgs, train_targets, test_targets, _, test_orig_targets = model_selection.train_test_split(
image_files, targets_enc, targets_orig, test_size=0.1, random_state=42)
train_dataset = dataset.ClassificationDataset(image_paths=train_imgs,
targets=train_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True)
test_dataset = dataset.ClassificationDataset(image_paths=test_imgs,
targets=test_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False)
model = CaptchaModel(num_chars=len(lbl_enc.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0,
patience=5,
verbose=True)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
valid_preds, valid_loss = engine.eval_fn(model, test_loader)
valid_cap_preds = []
for vp in valid_preds:
current_preds = decode_predictions(vp, lbl_enc)
valid_cap_preds.extend(current_preds)
print(list(zip(test_orig_targets, valid_cap_preds))[6:11])
print(
f"Epoch: {epoch}, train_loss={train_loss}, valid_loss={valid_loss}"
)
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
targets_orig = [x.split("/")[-1][:-4].split('_')[0] for x in image_files]
targets = [[c for c in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
lbl_enc = preprocessing.LabelEncoder()
lbl_enc.fit(targets_flat)
np.save(config.LABEL_ENCODER_SAVE_PATH, lbl_enc.classes_)
targets_enc = [lbl_enc.transform(x) for x in targets]
# print(targets_enc)
# new_targets_enc= []
# for i,target in enumerate(targets_enc):
# tmp = np.array([-1,-1,-1,-1,-1])
# for idx, item in enumerate(target):
# # print(idx)
# # print('i',i)
# tmp[idx] = item
# # print(image_files[i])
# new_targets_enc.append(tmp)
# print(new_targets_enc)
targets_enc = np.array(targets_enc)
targets_enc = targets_enc + 1
(
train_imgs,
test_imgs,
train_targets,
test_targets,
_,
test_targets_orig,
) = model_selection.train_test_split(
image_files, targets_enc, targets_orig, test_size=0.1, random_state=42
)
train_dataset = dataset.ClassificationDataset(
image_paths=train_imgs,
targets=train_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH),
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True,
)
test_dataset = dataset.ClassificationDataset(
image_paths=test_imgs,
targets=test_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH),
)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False,
)
model = CaptchaModel(num_chars=len(lbl_enc.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, factor=0.8, patience=5, verbose=True
)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
valid_preds, test_loss = engine.eval_fn(model, test_loader)
valid_captcha_preds = []
for vp in valid_preds:
current_preds = decode_predictions(vp, lbl_enc)
valid_captcha_preds.extend(current_preds)
combined = list(zip(test_targets_orig, valid_captcha_preds))
print(combined[:10])
test_dup_rem = test_targets_orig
accuracy = metrics.accuracy_score(test_dup_rem, valid_captcha_preds)
print(
f"Epoch={epoch}, Train Loss={train_loss}, Test Loss={test_loss} Accuracy={accuracy}"
)
scheduler.step(test_loss)
torch.save(model.state_dict(), config.MODEL_SAVE_PATH)
def run_training():
# Create pathlib.Path for the data
data_path = Path(config.data_dir)
image_files = list(data_path.glob("*.png"))
targets = []
targets_orig = []
targets_unique = set()
# Loop through each file and create target list
for file in data_path.iterdir():
targets_orig.append(file.stem) # append the filename
targets.append(list(file.stem)) # append the list of chars
targets_unique.update(list(file.stem)) # keep track of unique chars
msg = "Number of target data-points: {}, \nUnique chars: {} \n"
print(msg.format(len(targets), sorted(targets_unique)))
# Label encode
le = preprocessing.LabelEncoder()
le.fit(sorted(targets_unique))
targets_encoded = [le.transform(x) for x in targets]
targets_encoded = np.array(targets_encoded) + 1 # adding 1 because 0 represents "unkwown"
msg = "Encoded targets: \n{}"
print(msg.format(targets_encoded))
# Split the dataset
train_images, test_images, train_targets, test_targets, train_orig_targets, test_orig_targets = \
model_selection.train_test_split(
image_files, targets_encoded, targets_orig, test_size=0.1, random_state=42
)
train_dataset = dataset.ClassificationDataset(image_paths=train_images,
targets=train_targets,
resize=(config.image_height, config.image_width))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.batch_size,
num_workers=config.num_workers,
shuffle=True
)
test_dataset = dataset.ClassificationDataset(image_paths=test_images,
targets=test_targets,
resize=(config.image_height, config.image_width))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.batch_size,
num_workers=config.num_workers,
shuffle=False
)
# Create instance of the model and assign to gpu
model = CaptchaModel(num_chars=len(le.classes_))
# model.to(config.device)
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, factor=0.8, patience=5, verbose=True
)
prev_val_loss = sys.maxsize
for epoch in range(config.epochs):
# Train the model over all train batches
train_loss = train(model, train_loader, optimizer)
# Test the model over test batches
val_preds, val_loss = eval(model, test_loader)
# Print out the actual label and predicted labels
# Loop through and pass each batch to the decode function
val_preds_tmp = []
for vp in val_preds:
vp = decode_predictions(vp, le)
val_preds_tmp.extend(vp)
val_preds = val_preds_tmp
# Print out the first 5 predictions for the test set each epoch
print(f"Epoch: {epoch+1}, Train loss: {train_loss}, Val loss: {val_loss}")
pprint(list(zip(test_orig_targets, val_preds))[:5])
# Save the model if val_loss decreased
if val_loss <= prev_val_loss:
print(f"Val loss decreased from {prev_val_loss} to {val_loss}. Saving model.")
torch.save(model.state_dict(), Path(config.output_dir)/'captcha_model.pkl')
prev_val_loss = val_loss
print("\n\n")
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
# "/../../sdfrt.png" the next line only select the name of the file: sdfrt
targets_orig = [x.split("/")[-1][:-4] for x in image_files]
# sdfrt -> [s, d, f, r, t]
targets = [[c for c in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
lbl_enc = preprocessing.LabelEncoder()
lbl_enc.fit(targets_flat)
# Encode the targets
targets_enc = [lbl_enc.transform(x) for x in targets]
# Transform targets_enc to np.array
# The labels are encoded from 0 to N-1 where N is the number of labels
# we want to keep 0 to unknown so add 1
targets_enc = np.array(targets_enc) + 1
print(targets)
print(np.unique(targets_flat))
print(targets_enc)
print(len(lbl_enc.classes_))
# split in train, test for: imgs, targets, orig_targets
train_imgs, test_imgs, train_targets, test_targets, _, test_orig_targets = \
model_selection.train_test_split(image_files,
targets_enc, targets_orig,
test_size=0.1, random_state=42)
train_dataset = dataset.ClassificationDataset(image_paths=train_imgs, targets=train_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True
)
test_dataset = dataset.ClassificationDataset(image_paths=test_imgs, targets=test_targets,
resize=(config.IMAGE_HEIGHT, config.IMAGE_WIDTH))
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False,
)
model = CaptchaModel(num_chars=len(lbl_enc.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, factor=0.8, patience=5, verbose=True
)
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
valid_preds, valid_loss = engine.train_fn(model, test_loader)
valid_cap_preds = []
for valid_pred in valid_preds:
current_preds = decode_predictions(valid_pred, lbl_enc)
valid_cap_preds.extend(current_preds)
pprint(list(zip(test_orig_targets, valid_cap_preds))[6:10])
print(f"Epoch: {epoch}, train_loss={train_loss}, valid_loss={valid_loss}")
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
targets_orig = [
os.path.splitext(os.path.basename(i))[0] for i in image_files
]
targets = [[c for c in i] for i in targets_orig]
targets_flat = functools.reduce(operator.iconcat, targets, [])
# [j for i in targets for j in i ]
lbl_enc = preprocessing.LabelEncoder()
lbl_enc.fit(targets_flat)
target_enc = [lbl_enc.transform(i) for i in targets]
target_enc = np.array(target_enc) + 1
train_images, test_images, train_targets, test_targets, train_orig_targets, test_orig_targets = model_selection.train_test_split(
image_files, target_enc, targets_orig, test_size=0.1, random_state=42)
# print(train_images[0], test_images[0], train_targets[0], test_targets[0], train_orig_targets[0], test_orig_targets[0])
train_data = CaptchaImageDataset(image_paths=train_images,
targets=train_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
train_loader = data.DataLoader(train_data,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True)
test_data = CaptchaImageDataset(image_paths=test_images,
targets=test_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
test_loader = data.DataLoader(test_data,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False)
model = CaptchaModel(num_chars=len(lbl_enc.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=5,
verbose=True)
print('done')
losses = []
accuracy_scores = []
best_loss = 1e10
best_accuracy = 0
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model, train_loader, optimizer)
test_prediction, test_loss = engine.eval_fn(model, test_loader)
# print(test_prediction,test_loss)
# print(type(test_prediction[0]))
test_cap_prediction = []
for kk in test_prediction:
current_prediction = decode_predictions(kk, lbl_enc)
# print(current_prediction[0])
test_cap_prediction.extend(current_prediction)
# print(test_cap_prediction)
combined = list(zip(test_orig_targets, test_cap_prediction))
# print(combined[:10])
test_dup_rem = [remove_duplicates(c) for c in test_orig_targets]
# print(test_dup_rem)
accuracy = metrics.accuracy_score(test_dup_rem, test_cap_prediction)
accuracy_scores.append(accuracy)
losses.append(test_loss)
if best_loss > test_loss:
best_loss = test_loss
best_accuracy = accuracy
checkpoint = {
'epoch': epoch + 1,
'best_loss': test_loss,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_ckp(checkpoint, True, config.checkpoint_path,
config.best_model_path)
print('Saving Best Model')
print(
f"Epoch: {epoch}, train_loss:{train_loss}, test_loss:{test_loss}, Accuracy={accuracy}"
)
scheduler.step(test_loss)
print(
f"Accuracy of Model is {best_accuracy} and Test Loss is {best_loss,}")
fig = plt.figure()
ax = plt.axes()
ax.plot(accuracy_scores, label='Accuracy Scores')
plt.show()
fig = plt.figure()
ax = plt.axes()
ax.plot(losses, label='Loss Values')
plt.show()
def eval(model_dir,
data_dir,
batch_size=64,
log_dir='./logs',
use_gpu=True,
mode='captcha'):
"""
:param model_dir:
:param data_dir:
:param batch_size:
:param log_dir:
:param use_gpu:
:param mode:
:return:
"""
x_test, y_test = get_data_split(data_dir, modes=['test'])
if mode == 'captcha':
from model import CaptchaModel
elif mode == 'kaptcha':
from kaptcha_model import CaptchaModel
model = CaptchaModel()
gpu_available = torch.cuda.is_available()
if use_gpu and gpu_available:
model = model.cuda()
model_state = torch.load(model_dir)
else:
model_state = torch.load(model_dir,
map_location=lambda storage, loc: storage)
model.load_state_dict(model_state['network'])
test_ds = CaptchaLoader((x_test, y_test), shuffle=True)
test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=True)
model.eval()
acc_history = []
with tqdm(total=int(np.ceil(len(test_loader.dataset) / batch_size)),
desc='Eval') as eval_bar:
for _, (x, y) in enumerate(test_loader):
x = torch.tensor(x, requires_grad=False)
y = torch.tensor(y, requires_grad=False)
if use_gpu and gpu_available:
x = x.cuda()
y = y.cuda()
pred1, pred2, pred3, pred4 = model(x)
acc_mean = np.mean([
acc(pred1, y[:, 0]),
acc(pred2, y[:, 1]),
acc(pred3, y[:, 2]),
acc(pred4, y[:, 3])
])
pred = torch.stack((pred1, pred2, pred3, pred4), dim=-1)
multi_acc_mean = multi_acc(torch.argmax(pred, dim=1), y)
acc_history.append([acc_mean.item(), multi_acc_mean])
eval_bar.update()
eval_bar.set_postfix(acc=acc_mean, multi_acc=multi_acc_mean)
if not os.path.exists(log_dir):
os.mkdir(log_dir)
with open(os.path.join(log_dir, 'eval.json'), mode=r'w') as out_fp:
json.dump(acc_history, out_fp)
def run_training():
image_files = glob.glob(os.path.join(config.DATA_DIR, "*.png"))
image_files = image_files[:10]
print(f"Number of Images Found: {len(image_files)}")
# "../xywz.png" -> "xywz"
targets_orig = [x.split("/")[-1].split(".")[0] for x in image_files]
# separate the targets on character level
targets = [[char for char in x] for x in targets_orig]
targets_flat = [c for clist in targets for c in clist]
lbl_encoder = preprocessing.LabelEncoder()
lbl_encoder.fit(targets_flat)
targets_enc = [lbl_encoder.transform(x) for x in targets]
# label encodes from 0, so add 1 to start from 1: 0 will be saved for unknown
targets_enc = np.array(targets_enc) + 1
print(f"Number of Unique Classes: {len(lbl_encoder.classes_)}")
train_imgs, test_imgs, train_targets, test_targets, train_orig_targets, test_orig_targets = \
model_selection.train_test_split(image_files, targets_enc, targets_orig, test_size=0.1, random_state=42)
train_dataset = dataset.ClassificationDataset(image_paths=train_imgs,
targets=train_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=True)
test_dataset = dataset.ClassificationDataset(image_paths=test_imgs,
targets=test_targets,
resize=(config.IMAGE_HEIGHT,
config.IMAGE_WIDTH))
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
shuffle=False)
model = CaptchaModel(num_chars=len(lbl_encoder.classes_))
model.to(config.DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
factor=0.8,
patience=5,
verbose=True)
train_loss_data = []
test_loss_data = []
for epoch in range(config.EPOCHS):
train_loss = engine.train_fn(model,
train_loader,
optimizer,
save_model=True)
eval_preds, test_loss = engine.eval_fn(model, test_loader)
eval_captcha_preds = []
for vp in eval_preds:
current_preds = decode_predictions(vp, lbl_encoder)
eval_captcha_preds.extend(current_preds)
combined = list(zip(test_orig_targets, eval_captcha_preds))
pprint(combined[:10])
test_dup_rem = [remove_duplicates(c) for c in test_orig_targets]
accuracy = metrics.accuracy_score(test_dup_rem, eval_captcha_preds)
print(
f"Epoch={epoch}, Train Loss={train_loss}, Test Loss={test_loss} Accuracy={accuracy}"
)
scheduler.step(test_loss)
train_loss_data.append(train_loss)
test_loss_data.append(test_loss)
# print(train_dataset[0])
plot_loss(train_loss_data, test_loss_data, plot_path=config.PLOT_PATH)
print("done")