예제 소스
Qwen2.5-VL 차트 LoRA: PNG로 매수·매도·유지 학습하기 · 파일 59개
export_chart_pngs.py
"""
테스트 CSV(기본 cache/market_extreme_labels.csv)의 전체 봉을
``chart/{bar//10000:04d}/{bar:07d}.png`` 로 저장한다 (폴더당 최대 10000장).
지표는 ``RealTimeMarket.add_technical_indicators`` 와 같다.
차트는 ``ChartImageGenerator`` 가 그 시점까지의 최근 300봉을 한 장으로 그린다.
이미 있는 PNG 는 건너뛰므로 중단 후 다시 실행하면 이어서 그린다.
실행 (프로젝트 루트, 차트 전용 3.13 venv 권장 — LoRA/MLX 와 분리):
.venv-chart/bin/python export_chart_pngs.py --workers 1 --sleep-seconds 60
"""
from __future__ import annotations
import argparse
import logging
import os
import signal
import subprocess
import threading
import time
from collections import defaultdict
from multiprocessing import Pool
from pathlib import Path
import numpy as np
import pandas as pd
import talib
from ChartImageGenerator import (
ChartImageConfig,
ChartImageGenerator,
restart_kaleido_server,
stop_kaleido_server,
)
from lib.chart_png import chart_png_path, existing_chart_bars
logger = logging.getLogger(__name__)
TICKER = "KRW-ETH"
CHART_MAX_DATA = 300
RENDER_TIMEOUT_SEC = 45
def _process_tree() -> tuple[dict[int, list[int]], dict[int, str]]:
tree: dict[int, list[int]] = defaultdict(list)
cmds: dict[int, str] = {}
out = subprocess.check_output(["ps", "-axo", "pid=,ppid=,command="], text=True, errors="replace")
for line in out.splitlines():
parts = line.strip().split(None, 2)
if len(parts) < 3:
continue
try:
pid, ppid = int(parts[0]), int(parts[1])
except ValueError:
continue
tree[ppid].append(pid)
cmds[pid] = parts[2]
return tree, cmds
def _kill_headless_chrome_children() -> None:
"""kaleido 가 멈춘 headless Chrome 만 죽인다. 일반 Chrome 창은 건드리지 않는다."""
try:
tree, cmds = _process_tree()
except Exception:
return
stack = list(tree.get(os.getpid(), []))
descendants: list[int] = []
while stack:
pid = stack.pop()
descendants.append(pid)
stack.extend(tree.get(pid, []))
for pid in reversed(descendants):
cmd = cmds.get(pid, "")
kaleido_wrapper = "_unix_pipe_chromium_wrapper.py" in cmd
tmp_headless = "--headless" in cmd and "--user-data-dir=" in cmd and "/var/folders/" in cmd
if not (kaleido_wrapper or tmp_headless):
continue
try:
os.kill(pid, signal.SIGKILL)
except OSError:
pass
def _render_with_timeout(gen: ChartImageGenerator, df, i: int, max_data: int, scale: int) -> bytes:
done = threading.Event()
def _watch() -> None:
if not done.wait(RENDER_TIMEOUT_SEC):
_kill_headless_chrome_children()
watcher = threading.Thread(target=_watch, daemon=True)
watcher.start()
try:
return gen.render_image(df, idx=int(i) + 1, max_data=max_data, ticker=TICKER, scale=scale)
finally:
done.set()
watcher.join(timeout=1.0)
def _load_ohlcv(path: Path) -> pd.DataFrame:
usecols = ["bar_index", "open", "high", "low", "close", "volume"]
df = pd.read_csv(path, usecols=usecols)
if df.empty:
raise SystemExit(f"비어 있는 CSV: {path}")
df = df.sort_values("bar_index").reset_index(drop=True)
for c in ("open", "high", "low", "close", "volume"):
df[c] = pd.to_numeric(df[c], errors="coerce")
df = df.dropna(subset=["open", "high", "low", "close"]).reset_index(drop=True)
return df
def _with_realtime_indicators(df: pd.DataFrame) -> pd.DataFrame:
"""``RealTimeMarket.add_technical_indicators`` 와 같은 차트 지표.
RealTimeMarket 모듈은 Redis·pyupbit 를 import 하므로, 여기서는 그 계산식만 수행한다.
"""
out = df.copy()
close = np.asarray(out["close"], dtype=np.float64)
out["ma20"] = talib.SMA(close, 20)
out["ma60"] = talib.SMA(close, 60)
out["ma120"] = talib.SMA(close, 120)
upper, middle, lower = talib.BBANDS(close, timeperiod=40, nbdevup=2.3, nbdevdn=2.3, matype=0)
out["upper"] = upper
out["middle"] = middle
out["lower"] = lower
def _deg(series: pd.Series) -> pd.Series:
delta = series.diff()
return pd.Series(
np.degrees(np.arctan2(delta.fillna(0).to_numpy(dtype=np.float64), 3500.0)),
index=series.index,
)
out["ma20_deg"] = _deg(out["ma20"])
out["ma60_deg"] = _deg(out["ma60"])
out["ma120_deg"] = _deg(out["ma120"])
denom = (out["upper"] - out["lower"]).replace(0, np.nan)
out["PB"] = ((out["close"] - out["lower"]) / denom).replace([np.inf, -np.inf], np.nan).fillna(0.5).clip(0, 1)
if "MFI10" not in out.columns:
out["MFI10"] = 50.0
return out
def _render_shard(payload: tuple) -> tuple[int, int]:
df, positions, out_dir, max_data, scale, shard_id, sleep_sec = payload
out_dir = Path(out_dir)
pause = max(0.0, float(sleep_sec))
gen = ChartImageGenerator(ticker=TICKER, config=ChartImageConfig(show=False))
have = existing_chart_bars(out_dir)
written = 0
skipped = 0
t0 = time.perf_counter()
log = logging.getLogger(f"shard{shard_id}")
log.info("shard %s · 기존 PNG %s장", shard_id, len(have))
for n_done, i in enumerate(positions, start=1):
bar = int(df.iloc[i]["bar_index"])
dest = chart_png_path(out_dir, bar)
if bar in have:
skipped += 1
if skipped == 1 or skipped % 20000 == 0:
log.info("shard %s · 건너뜀 %s · bar=%s", shard_id, skipped, bar)
continue
png = b""
for attempt in range(4):
try:
png = _render_with_timeout(gen, df, int(i), max_data, scale)
break
except Exception as e:
log.warning("bar=%s 렌더 실패 %s/4: %s", bar, attempt + 1, e)
_kill_headless_chrome_children()
restart_kaleido_server(timeout=RENDER_TIMEOUT_SEC)
gen = ChartImageGenerator(ticker=TICKER, config=ChartImageConfig(show=False))
time.sleep(2.0 * (attempt + 1))
else:
raise RuntimeError(f"bar={bar} PNG 렌더를 4회 실패")
dest.parent.mkdir(parents=True, exist_ok=True)
dest.write_bytes(png)
have.add(bar)
written += 1
elapsed = time.perf_counter() - t0
log.info(
"shard %s · 저장 %s · 건너뜀 %s · bar=%s · %.4f장/초 · %s/%s",
shard_id,
written,
skipped,
bar,
written / max(elapsed, 1e-6),
n_done,
len(positions),
)
if pause > 0 and n_done < len(positions):
stop_kaleido_server()
_kill_headless_chrome_children()
log.info("CPU 풀 양보 · %.0f초 대기 (다음 장)", pause)
time.sleep(pause)
return written, skipped
def export_pngs(
df: pd.DataFrame,
out_dir: Path,
*,
max_data: int,
scale: int,
start_bar: int,
workers: int,
sleep_sec: float,
) -> tuple[int, int]:
out_dir.mkdir(parents=True, exist_ok=True)
positions = np.flatnonzero(df["bar_index"].to_numpy() >= int(start_bar))
if len(positions) == 0:
return 0, 0
pause = max(0.0, float(sleep_sec))
n_workers = max(1, int(workers))
if pause > 0:
n_workers = 1
shards = [arr.tolist() for arr in np.array_split(positions, n_workers) if len(arr)]
payloads = [
(df, shard, str(out_dir), int(max_data), int(scale), i, pause)
for i, shard in enumerate(shards)
]
if n_workers == 1:
return _render_shard(payloads[0])
written = skipped = 0
with Pool(processes=len(payloads)) as pool:
for w, s in pool.imap_unordered(_render_shard, payloads):
written += w
skipped += s
logger.info("누적 저장 %s · 건너뜀 %s", written, skipped)
return written, skipped
def main() -> None:
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
parser = argparse.ArgumentParser(description="테스트 CSV 전체 봉을 chart PNG 로 저장")
parser.add_argument("--csv", type=Path, default=Path("cache/market_extreme_labels.csv"))
parser.add_argument("--out", type=Path, default=Path("chart"))
parser.add_argument("--max-data", type=int, default=CHART_MAX_DATA)
parser.add_argument("--scale", type=int, default=1)
parser.add_argument(
"--start-bar",
type=int,
default=300,
help="이 bar_index 부터 저장 (ma120·300봉 창이 채워진 뒤)",
)
parser.add_argument("--workers", type=int, default=4, help="동시에 그릴 프로세스 수")
parser.add_argument(
"--sleep-seconds",
type=float,
default=60.0,
help="한 장 저장 후 CPU 풀을 비우기 위해 쉬는 초 (0 이면 쉬지 않음)",
)
args = parser.parse_args()
if not args.csv.is_file():
raise SystemExit(f"CSV 없음: {args.csv}")
logger.info("CSV 로드: %s", args.csv)
raw = _load_ohlcv(args.csv)
logger.info("지표 계산 (RealTimeMarket.add_technical_indicators) · %s행", len(raw))
df = _with_realtime_indicators(raw)
need = ["ma20", "ma60", "ma120", "ma20_deg", "upper", "lower", "volume"]
missing = [c for c in need if c not in df.columns]
if missing:
raise SystemExit(f"차트 컬럼 없음: {missing}")
logger.info(
"PNG 저장 시작 → %s (start_bar=%s · sleep=%.0fs)",
args.out,
args.start_bar,
max(0.0, float(args.sleep_seconds)),
)
written, skipped = export_pngs(
df,
args.out,
max_data=max(1, int(args.max_data)),
scale=max(1, int(args.scale)),
start_bar=int(args.start_bar),
workers=int(args.workers),
sleep_sec=float(args.sleep_seconds),
)
logger.info("완료 · 새로 저장 %s · 기존 건너뜀 %s", written, skipped)
if __name__ == "__main__":
main()