예제 소스
차트 REINFORCE: 규칙 교사와 KL로 정책망 준비하기 · 파일 59개
ChartImageGenerator.py
import logging
from dataclasses import dataclass
from io import BytesIO
from typing import List, Optional, Sequence, Tuple
logger = logging.getLogger(__name__)
_KALEIDO_SERVER_STARTED = False
def ensure_kaleido_server(*, timeout: int = 40) -> None:
"""Chrome 을 장마다 켜지 않고 kaleido 서버 1개를 재사용한다."""
global _KALEIDO_SERVER_STARTED
try:
import kaleido
except ImportError:
return
try:
kaleido.start_sync_server(silence_warnings=True, n=1, timeout=timeout)
_KALEIDO_SERVER_STARTED = True
except Exception as e:
logger.warning("kaleido 서버 시작 실패: %s", e)
def stop_kaleido_server() -> None:
try:
import kaleido
kaleido.stop_sync_server(silence_warnings=True)
except Exception:
pass
global _KALEIDO_SERVER_STARTED
_KALEIDO_SERVER_STARTED = False
def restart_kaleido_server(*, timeout: int = 40) -> None:
stop_kaleido_server()
ensure_kaleido_server(timeout=timeout)
@dataclass
class ChartImageConfig:
rows: int = 3
row_heights: Sequence[float] = (0.5, 0.2, 0.3)
showlegend: bool = False
show: bool = False
template: str = "seaborn"
class ChartImageGenerator:
"""
Plotly 기반 차트(캔들+MA+ma20_deg+거래 마커)를 이미지/figure로 생성하는 클래스.
- 기존 get_fig(df, idx, max_data) 로직을 그대로 클래스 메서드로 옮겼습니다.
- 현재는 "한 화면(윈도우) 1장"만 생성합니다. (추후 세로 분할 확장 가능)
"""
def __init__(
self,
ticker: Optional[str] = None,
config: Optional[ChartImageConfig] = None,
) -> None:
self.ticker = ticker
self.config = config or ChartImageConfig()
# lazy import: plotly/kaleido 미설치 환경에서 import-time 크래시 방지
self._go = None
self._make_subplots = None
def _ensure_plotly(self) -> None:
if self._go is not None and self._make_subplots is not None:
return
try:
import plotly.graph_objects as go # type: ignore
from plotly.subplots import make_subplots # type: ignore
except ImportError as e:
raise ImportError(
"plotly가 필요합니다. `pip install plotly` 후 다시 실행해 주세요."
) from e
self._go = go
self._make_subplots = make_subplots
@staticmethod
def _require_columns(df, cols: Sequence[str], strict: bool = True) -> None:
missing = [c for c in cols if c not in df.columns]
if missing and strict:
raise KeyError(f"df에 필요한 컬럼이 없습니다: {missing}")
def build_figure(self, df, idx: int, max_data: int = 300, ticker: Optional[str] = None):
"""
기존 get_fig 로직과 동일하게 plotly figure 생성.
"""
self._ensure_plotly()
go = self._go
ms = self._make_subplots
if ticker is None:
ticker = self.ticker or ""
# 필요한 데이터만 추출
df = df.head(idx).tail(max_data)
df = df.reset_index(drop=True)
# 최소한 아래 컬럼들은 기존 get_fig에서 사용됨
required_cols = [
"open",
"high",
"low",
"close",
"upper",
"ma20",
"ma60",
"ma120",
"lower",
"ma20_deg",
"volume",
"PB",
"MFI10",
]
self._require_columns(df, required_cols, strict=False)
x = df.index
candle = go.Candlestick(
x=x,
open=df["open"],
high=df["high"],
low=df["low"],
close=df["close"],
increasing_line_color="red",
decreasing_line_color="blue",
showlegend=False,
)
ma20 = go.Scatter(
x=x,
y=df["ma20"],
line=dict(color="red", width=2),
name="ma20",
showlegend=False,
)
ma60 = go.Scatter(
x=x,
y=df["ma60"],
line=dict(color="blue", width=2),
name="ma60",
showlegend=False,
)
ma120 = go.Scatter(
x=x,
y=df["ma120"],
line=dict(color="black", width=2),
name="ma120",
showlegend=False,
)
ma20_deg = go.Scatter(
x=x,
y=df["ma20_deg"],
fill="tozeroy",
fillcolor="rgba(255, 0, 0, 0.3)",
line=dict(color="red", width=2),
name="ma20_deg",
showlegend=False,
)
volume = go.Bar(
x=x,
y=df["volume"],
marker_color="red",
name="volume",
showlegend=False,
)
# 아래는 원본 코드에서 만들었지만 add_trace를 하지 않았던 것들.
# (향후 필요하면 traces 옵션으로 추가 확장하기 좋게 유지)
_ = (
go.Scatter,
go.Bar,
)
fig = ms(
rows=3,
cols=1,
start_cell="top-left",
row_heights=list(self.config.row_heights),
vertical_spacing=0.0,
)
fig.add_trace(candle, row=1, col=1)
fig.add_trace(ma20, row=1, col=1)
fig.add_trace(ma60, row=1, col=1)
fig.add_trace(ma120, row=1, col=1)
fig.add_trace(ma20_deg, row=2, col=1)
fig.add_trace(volume, row=3, col=1)
# 추세추종/역추세 전략 표시 (triangle 마커)
# 원본 코드에서는 조건 미충족 시 trend_fol 미정의 상태 가능성이 있어,
# 여기서는 조건 충족일 때만 add_trace 합니다.
for i in df.index:
trend_fol = None
if "PB" in df.columns and "MFI10" in df.columns:
if df["PB"][i] > 0.8 and df["MFI10"][i] > 80:
trend_fol = go.Scatter(
x=[df.index[i]],
y=[df["close"][i]],
marker_color="orange",
marker_size=20,
marker_symbol="triangle-up",
opacity=0.7,
showlegend=False,
)
elif df["PB"][i] < 0.2 and df["MFI10"][i] < 20:
trend_fol = go.Scatter(
x=[df.index[i]],
y=[df["close"][i]],
marker_color="darkblue",
marker_size=20,
marker_symbol="triangle-down",
opacity=0.7,
showlegend=False,
)
if trend_fol is not None:
fig.add_trace(trend_fol, row=1, col=1)
# title / axis 스타일은 원본 코드의 흐름을 유지
fig.update_layout(
autosize=True,
xaxis1_rangeslider_visible=False,
xaxis2_rangeslider_visible=False,
# 이미지 학습용: 불필요한 외곽 여백 제거
margin=dict(l=0, r=0, t=0, b=0, pad=0),
template=self.config.template,
)
fig.update_xaxes(
zeroline=True,
zerolinewidth=1,
zerolinecolor="black",
showgrid=True,
gridwidth=2,
gridcolor="lightgray",
showline=True,
linewidth=2,
linecolor="black",
mirror=True,
showticklabels=False,
)
fig.update_yaxes(
tickformat=",d",
zeroline=True,
zerolinewidth=1,
zerolinecolor="black",
showgrid=True,
gridwidth=2,
gridcolor="lightgray",
showline=True,
linewidth=2,
linecolor="black",
mirror=True,
showticklabels=False,
)
if self.config.show:
fig.show()
return fig
def render_image(
self,
df,
idx: int,
max_data: int = 300,
ticker: Optional[str] = None,
output_format: str = "png",
scale: int = 2,
) -> bytes:
"""
figure를 이미지 바이트로 반환 (plotly.kaleido 필요).
"""
fig = self.build_figure(df, idx=idx, max_data=max_data, ticker=ticker)
ensure_kaleido_server()
try:
img_bytes = fig.to_image(format=output_format, scale=scale)
except Exception as e:
msg = str(e) or repr(e)
if "requires Google Chrome to be installed" in msg:
raise RuntimeError(
"이미지 변환 실패: kaleido는 Chrome/Chromium이 필요합니다.\n"
"- Ubuntu라면 예: `sudo apt install -y chromium-browser` 또는 `sudo apt install -y chromium`\n"
"- 또는 Google Chrome을 설치하세요.\n"
"설치 후 동일한 venv에서 다시 실행해 주세요."
) from e
if "BrowserFailedError" in msg or "The browser seemed to close immediately" in msg:
raise RuntimeError(
"이미지 변환 실패: kaleido가 브라우저(Chrome/Chromium)를 실행했지만 즉시 종료되었습니다.\n"
"- 현재 시도한 브라우저 경로가 로그에 나오면(예: `/snap/bin/chromium`) 그 Chromium이 동작하는지 먼저 확인하세요.\n"
"- Snap chromium에서 자주 발생하면 apt chromium / google-chrome으로 변경하는 게 보통 해결됩니다.\n"
"- kaleido 1.x는 내부적으로 브라우저가 필요합니다."
) from e
raise RuntimeError(
"이미지 변환에 실패했습니다. kaleido가 필요할 수 있습니다. "
"`pip install -U kaleido` 후 다시 시도해 주세요."
) from e
return img_bytes
@staticmethod
def slice_image_by_x(
img_bytes: bytes,
*,
n_slices: int,
overlap_px: int = 0,
output_format: str = "PNG",
) -> List[bytes]:
"""
렌더된 차트 이미지(bytes)를 x축(가로) 방향으로 잘라 여러 장(bytes)으로 반환.
- **n_slices**: 가로를 몇 장으로 자를지 (>= 1)
- **overlap_px**: 조각 간 겹침 픽셀 (기본 0)
- **output_format**: 반환 이미지 포맷 (기본 PNG)
Pillow(PIL)가 필요합니다: `pip install pillow`
"""
if n_slices < 1:
raise ValueError("n_slices는 1 이상이어야 합니다.")
if overlap_px < 0:
raise ValueError("overlap_px는 0 이상이어야 합니다.")
try:
from PIL import Image # type: ignore
except ImportError as e:
raise ImportError(
"이미지 자르기는 Pillow가 필요합니다. `pip install pillow` 후 다시 시도해 주세요."
) from e
img = Image.open(BytesIO(img_bytes))
img = img.convert("RGBA")
w, h = img.size
base = w // n_slices
remainder = w % n_slices
chunks: List[bytes] = []
x0 = 0
for i in range(n_slices):
seg_w = base + (1 if i < remainder else 0)
left = max(0, x0 - overlap_px)
right = min(w, x0 + seg_w + overlap_px)
crop = img.crop((left, 0, right, h))
buf = BytesIO()
crop.save(buf, format=output_format)
chunks.append(buf.getvalue())
x0 += seg_w
return chunks
def render_and_slice_by_x(
self,
df,
idx: int,
*,
max_data: int = 300,
ticker: Optional[str] = None,
n_slices: int = 3,
overlap_px: int = 0,
scale: int = 2,
) -> Tuple[bytes, List[bytes]]:
"""
df로 차트 이미지를 1장 렌더링한 뒤, x축 방향으로 잘라 반환.
반환: (원본_이미지_bytes, [조각1, 조각2, ...])
"""
base_img = self.render_image(
df=df,
idx=idx,
max_data=max_data,
ticker=ticker,
output_format="png",
scale=scale,
)
parts = self.slice_image_by_x(
base_img,
n_slices=n_slices,
overlap_px=overlap_px,
output_format="PNG",
)
return base_img, parts