[파이썬] 주식 Chart Image 날짜 지정 저장

주식 Chart를 Image를 받을 수 있는 코드입니다.

주어진 날짜로 Train(UP/DOWN), Validation(UP/DOWN), Test(UP/DOWN) 각각의 폴더를 만들고 Image를 다운로드합니다.

CNN 등을 처리할 때 프로그램 내에 Image 다운코드까지 만들 수 있겠지만 Image 다운 시간이 많이 걸리기 때문에 Image 데이터는 따로 받아두는 것은 좋은 것 같습니다.

#import matplotlib.pyplot as plt
#from numpy import sqrt
import os
import shutil
import pandas_datareader as web
import random
import numpy as np
import pandas as pd
import datetime
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from mplfinance.original_flavor import candlestick2_ohlc
from pykrx import stock

STOCK_CODE = ['069500','005930','051910','035420','005380','035720','068270','066570','005490','051900','017670']
STOCK_NAME = ['KODEX200','삼성전자','LG화학','네이버','현대차','카카오','셀트리온','LG전자','POSCO','LG생활건강','SK텔레콤']
TRAIN_START = '2019-01-01'
VALI_START = '2020-12-31'
VALI_END = '2021-12-31'
TEST_END = '2022-08-16'

CODE_IDX = 5

#fileName = 'STOCK_DATA' + '_' + START + '_' + END + '.txt'
train_dir_base = './stockData2/train'
vali_dir_base = './stockData2/vali'
test_dir_base = './stockData2/test'
train_dir_up = train_dir_base + '/up'
train_dir_down = train_dir_base + '/down'
vali_dir_up = vali_dir_base + '/up'
vali_dir_down = vali_dir_base + '/down'
test_dir_up = test_dir_base + '/up'
test_dir_down = test_dir_base + '/down'


for i in range(6):          # 각 폴더별 삭제 및 생성
    s_dir = ''
    if i == 0:
        s_dir = train_dir_up
    elif i == 1:
        s_dir = train_dir_down
    elif i == 2:
        s_dir = vali_dir_up
    elif i == 3:
        s_dir = vali_dir_down
    elif i == 4:
        s_dir = test_dir_up
    elif i == 5:
        s_dir = test_dir_down

    if not os.path.exists(s_dir) == False:
        print('##### delete ' + s_dir + ' directory #####')
        shutil.rmtree(s_dir)

    print('##### make ' + s_dir + ' directory #####')
    os.makedirs(s_dir)


#oLog = open(fileName, 'w')
now = datetime.datetime.now()
nowDatetime = now.strftime('%Y%m%d_%H%M')

# df = web.DataReader(STOCK_CODE[CODE_IDX]+'.KS', data_source='yahoo', start=TRAIN_START, end=TEST_END)
df = stock.get_market_ohlcv(TRAIN_START, TEST_END, STOCK_CODE[CODE_IDX])        # 1 개 주식 데이터만 사용
df.columns = ['Open', 'High', 'Low', 'Close', 'Volume']
df['ma5'] = df['Close'].rolling(window=5).mean()
df['ma20'] = df['Close'].rolling(window=20).mean()
df['ma60'] = df['Close'].rolling(window=60).mean()
df['ma120'] = df['Close'].rolling(window=120).mean()
df['date'] = df.index.astype(str).str[0:10]
# 볼린전밴드
w = 20
k = 2
df['bol_upper'] = df['Close'].rolling(window=w).mean() + (k * df['Close'].rolling(window=w).std())
df['bol_down'] = df['Close'].rolling(window=w).mean() - (k * df['Close'].rolling(window=w).std())

df.dropna()

# train, validation, test 데이터 분리(인덱스 찾기)
train_start_idx = -1
train_end_idx = -1
vali_start_idx = -1
vali_end_idx = -1
test_start_idx = -1
test_end_idx = -1
date = df.date.tolist()
for i in range(len(date)):
    if train_start_idx == -1 and (date[i] == TRAIN_START or date[i] > TRAIN_START):
        train_start_idx = i
    elif train_end_idx == -1 and (date[i] == VALI_START or date[i] > VALI_START):
        train_end_idx = i
        vali_start_idx = i
    elif vali_end_idx == -1 and (date[i] == VALI_END or date[i] > VALI_END):
        vali_end_idx = i
        test_start_idx = i
    elif test_end_idx == -1 and (date[i] == TEST_END or date[i] > TEST_END or i == len(date)-1):
        test_end_idx = i

# 데이터 분리
X_train = df.iloc[:train_end_idx+2]             # +1은 -1까지 end row가 되기때문이고 +1은 내일의 up,down을 구분하기 위한 값
X_vali = df.iloc[vali_start_idx:vali_end_idx+2]
X_test = df.iloc[test_start_idx:]

day = 5     # Image에 보여줄 데이터 일수
for t in range(3):    # Train, Vali, Test 구분해서 처리
    data = pd.DataFrame()
    dir_base = ''
    if t == 0:        # Train, Vali, Test에 따른 데이터 및 디렉토리 지정
        data = X_train
        dir_base = train_dir_base
    elif t == 1:
        data = X_vali
        dir_base = vali_dir_base
    else:
        data = X_test
        dir_base = test_dir_base

    date_idx = data['date']
    for i in range(day,len(data)-1):   # 데이터 Row 만큼 처리
        s = i - day
        e = i + 1
        view = data.iloc[s:e]
        index = view.date

        target = ''
        target_dir = './'
        #if i < len(df)-1:
        if data['Close'][i+1] > data['Close'][i]:
            target_dir = dir_base + '/up'
            target = 'up'
            # text = '▲'
        else:
            target_dir = dir_base + '/down'
            target = 'down'
            # text = '▲'          # '▼'

        # font01 = {'family': 'NanumBarunGothic', 'color': 'red', 'weight': 'bold', 'size': 120}
        # font02 = {'family': 'NanumBarunGothic', 'color': 'blue', 'weight': 'bold', 'size': 120}

        fig = plt.figure(figsize=(15, 7))   # figsize, DPI로 image 크기 결정
        ax_main = fig.add_subplot(1, 1, 1)
        #ax_main.xaxis.set_major_locator(ticker.MaxNLocator(10))
        #ax_main.xaxis.set_major_formatter(ticker.FuncFormatter(x_date))
        # ax = plt.gca()
        # ax.axes.xaxis.set_visible(False)
        # ax.axes.yaxis.set_visible(False)
        # ax.axis('off')
        ax_main.plot(index, view['ma5'], linewidth=12)      # linewidth로 그래프 Line의 두께 조절
        ax_main.plot(index, view['ma20'], linewidth=12)
        ax_main.plot(index, view['ma60'], linewidth=12)
        ax_main.plot(index, view['ma120'], linewidth=12)
        ax_main.plot(index, view['bol_upper'], linewidth=12)
        ax_main.plot(index, view['bol_down'], linewidth=12)
        candlestick2_ohlc(ax_main, view['Open'], view['High'], view['Low'], view['Close'], width=0.5, colorup='r', colordown='b',alpha=1)
        #plt.tight_layout()
        ax_main.xaxis.set_major_locator(ticker.MaxNLocator(7))
        ax_main.grid(color='k',linewidth=1)

        y = int(view['Close'][day])
        # if target == 'up':
        #     ax_main.text(index[-3], y, text, fontdict=font01)
        # else:
        #     ax_main.text(index[-3], y, text, fontdict=font02)
        #plt.show()
        img_name = STOCK_CODE[CODE_IDX] + '_' + date_idx[i][0:4] + date_idx[i][5:7] + date_idx[i][8:10] + '_' + target + '.png'
        plt.savefig(target_dir + '/' + img_name, dpi=4)
        plt.close(fig)
        print('[',i+1,'/',len(data),'] target=',target_dir,' / name=',img_name)

train, vali, test 폴더를 만들고 Image 를 다운 받았습니다.