load_yes_no_helmet.py

def load_data(path):
    x_train = []
    y_train = []
    yes_helmet_location=os.path.join(location,"yes_helmet")
    no_helmet_location=os.path.join(location,"no_helmet")
    yes_helmet_files=os.listdir(yes_helmet_location)
    no_helmet_files=os.listdir(no_helmet_location)
    
    for ffile in yes_helmet_files[:5000]:
        img=cv2.imread(os.path.join(yes_helmet_location,ffile))
        tr_x = image.img_to_array(img)
        tr_x = preprocess_input(tr_x)
        x_train.append(img)
        y_train.append(1)
    for ffile in no_helmet_files:
        img=cv2.imread(os.path.join(no_helmet_location,ffile))
        tr_x = image.img_to_array(img)
        tr_x = preprocess_input(tr_x)
        x_train.append(img)
        y_train.append(0)
    return np.array(x_train), to_categorical(y_train)

location="raw_yolo_keras/train/"
x_train, y_train = load_data(location)

print(type(y_train))
print(y_train.shape)    # 808,4
print(x_train.shape)    # 808,299,299,3

x_train, x_test, y_train, y_test = train_test_split(x_train, y_train, test_size=0.33, random_state=42)
print(x_train.shape,x_test.shape)
print(y_train.shape,y_test.shape)