# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import cv2 as cv
import os
from matplotlib import pyplot as plt
import numpy as np
llist = {'0':0,'1':0,'2':0,'3':0,'4':0,'5':0,'6':0,'7':0,'8':0,'9':0,
         'a':0,'b':0,'c':0,'d':0,'e':0,'f':0,'g':0,'h':0,'i':0,'j':0,'k':0,'l':0,'m':0,'n':0,
         'o':0,'p':0,'q':0,'r':0,'s':0,'t':0,'u':0,'v':0,'w':0,'x':0,'y':0,'z':0,
        'A':0,'B':0,'C':0,'D':0,'E':0,'F':0,'G':0,'H':0,'I':0,'J':0,'K':0,'L':0,'M':0,
         'N':0,'O':0,'P':0,'Q':0,'R':0,'S':0,'T':0,'U':0,'V':0,'W':0,'X':0,'Y':0,'Z':0}

data_path = '../data/color_sample'
file_names = os.listdir(data_path)
file_names
# -

llist['a']

SAMPLE1_PATH = data_path+'/'+file_names[0]
SAMPLE2_PATH = data_path+'/'+file_names[1]
SAMPLE3_PATH = data_path+'/'+file_names[2]
SAMPLE4_PATH = data_path+'/'+file_names[3]
SAMPLE5_PATH = data_path+'/'+file_names[4]
SAMPLE6_PATH = data_path+'/'+file_names[5]
SAMPLE7_PATH = data_path+'/'+file_names[6]
SAMPLE8_PATH = data_path+'/'+file_names[7] 

img = cv.imread(SAMPLE6_PATH)
plt.imshow(img)
plt.show()
n_img = np.zeros((img.shape[0],img.shape[1]))
img_aft = cv.normalize(img, n_img, 0,255,cv.NORM_MINMAX)
plt.imshow(img_aft)
plt.show()

gray = cv.cvtColor(img_aft,cv.COLOR_BGR2GRAY)
plt.imshow(gray)
plt.show()
ret, thresh = cv.threshold(gray,0,255,cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
#ret, thresh = cv.threshold(gray,0,255,cv.THRESH_OTSU)
plt.imshow(thresh)
plt.show()

import copy
im2,contours,hierarchy = cv.findContours(thresh,cv.RETR_EXTERNAL,cv.CHAIN_APPROX_SIMPLE)
filter_containor = []
temp_img = copy.deepcopy(img)
for i in range(0,len(contours)):  
    x, y, w, h = cv.boundingRect(contours[i])   
    newimage=img[y:y+h,x:x+w] # 先用y确定高，再用x确定宽
    nrootdir=("cut_image/")
    if h<5 and w<5:
        continue
    filter_containor.append([x, y, w, h])
    cv.rectangle(temp_img, (x,y), (x+w,y+h), (153,153,0), 1)
    if not os.path.isdir(nrootdir):
        os.makedirs(nrootdir)
    cv.imwrite( nrootdir+str(i)+".jpg",newimage) 
    print (x, y, w, h)
plt.imshow(temp_img)
plt.show()


def seg_char(img_folder, img_name, save_folder):
    img = cv.imread(img_folder+'/'+img_name)
    name_list = img_name.split('_')[0]
    gray = cv.cvtColor(img_aft,cv.COLOR_BGR2GRAY)
    ret, thresh = cv.threshold(gray,0,255,cv.THRESH_BINARY_INV+cv.THRESH_OTSU)

    n_img = np.zeros((img.shape[0],img.shape[1]))
    img_aft = cv.normalize(thresh, n_img, 0,255,cv.NORM_MINMAX)
    im2,contours,hierarchy = cv.findContours(thresh,cv.RETR_LIST,cv.CHAIN_APPROX_SIMPLE)
    filter_containor = []
    temp_img = copy.deepcopy(img)
    for i in range(0,len(contours)):  
        x, y, w, h = cv.boundingRect(contours[i])   
        newimage=img[y:y+h,x:x+w] # 先用y确定高，再用x确定宽
        char_class = name_list[-1-i]
        #nrootdir=save_folder#("cut_image/")
        if h<5 and w<5:
            continue
        filter_containor.append([x, y, w, h])
        cv.rectangle(temp_img, (x-2,y-2), (x+w+2,y+h+2), (153,153,0), 1)
        file_path = save_folder+char_class+'/'
        if not os.path.isdir(file_path):
            os.makedirs(file_path)
        cv.imwrite( file_path+str(llist[char_class])+".png",newimage)
        llist[char_class] +=1
        print (x, y, w, h)
    plt.imshow(temp_img)
    plt.show()


seg_char(data_path, file_names[5], "test/" )

# +
import sys, os
sys.path.insert(1, '/home/ning_a/Desktop/CAPTCHA/base_solver/base_solver_char')
import cv2 as cv
from matplotlib import pyplot as plt
import numpy as np
import operator
import copy
from captcha_cnn_model import CNN, Generator
import torch
from torch.autograd import Variable
import captcha_setting
import my_dataset
from torchvision.utils import save_image
import cv2 as cv
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
from PIL import Image
import copy
import operator
import torch.nn as nn
import captcha_setting
from matplotlib import cm

llist = {'0':0,'1':0,'2':0,'3':0,'4':0,'5':0,'6':0,'7':0,'8':0,'9':0,
         'a':0,'b':0,'c':0,'d':0,'e':0,'f':0,'g':0,'h':0,'i':0,'j':0,'k':0,'l':0,'m':0,'n':0,
         'o':0,'p':0,'q':0,'r':0,'s':0,'t':0,'u':0,'v':0,'w':0,'x':0,'y':0,'z':0,
        'A':0,'B':0,'C':0,'D':0,'E':0,'F':0,'G':0,'H':0,'I':0,'J':0,'K':0,'L':0,'M':0,
         'N':0,'O':0,'P':0,'Q':0,'R':0,'S':0,'T':0,'U':0,'V':0,'W':0,'X':0,'Y':0,'Z':0}
generator = Generator()
generator.load_state_dict(torch.load('/home/ning_a/Desktop/CAPTCHA/base_solver/base_solver_char/7800.pkl'))
generator.eval()
# data_path ='/home/ning_a/Desktop/CAPTCHA/test_generator/test'# '../data/color_sample'
data_path = '/home/ning_a/Desktop/CAPTCHA/dataset_darkweb_bias/kaggle_captcha_4_letter/train/'
save_path = '/home/ning_a/Desktop/CAPTCHA/dataset_darkweb_bias/kaggle_captcha_4_letter/train_char_gen/'
transform = transforms.Compose([
    # transforms.ColorJitter(),
    transforms.Grayscale(),
    transforms.ToTensor(),
    # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
file_names = os.listdir(data_path)
# file_names
def seg_char(img_folder, img_name, save_folder):
    img = cv.imread(img_folder+'/'+img_name)
    torch_img = Image.open(img_folder+'/'+img_name)
#     print(img.shape)
    torch_img = torch_img.resize((160, 60))
#     torch_img = cv.resize(torch_img,(160, 60), interpolation = cv.INTER_CUBIC)
#     print(torch_img.shape)
    
    torch_img = transform(torch_img)
    gen_img = generator(torch_img)
    gen_img2 = gen_img.data.cpu().numpy()
    target_img = gen_img2[0][0]
    target_img = target_img*255
    cv.imwrite( "temp.jpg",target_img) 
    gen_img = cv.imread( "temp.jpg")
    gen_img = torch_img = cv.resize(gen_img,(72, 24), interpolation = cv.INTER_CUBIC)
    name_list = img_name.split('.')[0]
#     print(name_list)
    gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
    gray_2 = cv.cvtColor(gen_img,cv.COLOR_BGR2GRAY)
    ret, thresh = cv.threshold(gray,0,255,cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
    
    
    n_img = np.zeros((img.shape[0],img.shape[1]))
    img_aft = cv.normalize(thresh, n_img, 0,255,cv.NORM_MINMAX)
    im2,contours,hierarchy = cv.findContours(thresh,cv.RETR_EXTERNAL,cv.CHAIN_APPROX_TC89_KCOS)
    filter_containor = []
    temp_img = copy.deepcopy(img)
#     print("seg #: ",len(contours))
    num_index = 0
    if(len(contours)!=4):
        raise Exception("Sorry, no numbers below zero") 
    for i in range(0,len(contours)): 
        x, y, w, h = cv.boundingRect(contours[i])
        
        filter_containor.append([x, y, w, h])
#             continue
#         if h<35 or w<5:
#             continue
#         filter_containor.append([x, y, w, h])
    # sort contour list
#     print(filter_containor)
    filter_containor = sorted(filter_containor, key=operator.itemgetter(0))
    
    for i in range(0,len(filter_containor)):  
        #x, y, w, h = cv.boundingRect(contours[i])
        #try:
#         print(name_list, i)
        char_class = name_list[i]
        x = max(0, filter_containor[i][0]-2)
        y = max(0,filter_containor[i][1]-2)
        if(char_class=='i'):
            y = 0# max(0,filter_containor[i][1]-2)
        xr = min(159,filter_containor[i][0]+filter_containor[i][2]+2 )
        if(char_class=='i'):
            yb = 59# max(0,filter_containor[i][1]-2)
        yb = min(59, filter_containor[i][1]+filter_containor[i][3]+2 )
#         newimage=gray[y:yb,x:xr] # 先用y确定高，再用x确定宽
        newimage=gray_2[y:yb,x:xr]
        #char_class = name_list[-1-i]
        #nrootdir=save_folder#("cut_image/")
        
        num_index+=1

        cv.rectangle(temp_img, (x,y), (xr,yb), (153,153,0), 1)
        file_path = save_folder+char_class+'/'
        if not os.path.isdir(file_path):
            os.makedirs(file_path)
        newimage = cv.resize(newimage,(30, 60), interpolation = cv.INTER_CUBIC)
#         print(char_class)
        cv.imwrite( file_path+(str(llist[char_class]).upper())+".png",newimage)
        llist[char_class] +=1
#         except:
#             print(img_name)
#             plt.imshow(temp_img)
#             plt.show()
        #print (x, y, xr, h)
#     if('i' in name_list):
#         print(filter_containor)
#         plt.imshow(temp_img)
#         plt.show()
total = 0
correct = 0
for each_image in file_names:
    #print(each_image)
    total+=1
    try:
        seg_char(data_path, each_image, save_path )
        correct += 1
    except:
        continue
print(total, correct, float(correct)/float(total))
# -



StarcoderdataJupyterScriptsDedupFiltered