예제 소스

test_img.py

import argparse
import os
from datetime import datetime, timedelta
from io import StringIO
import sys

import numpy as np
import pandas as pd

from ChartImageGenerator import ChartImageConfig, ChartImageGenerator


def _dummy_df(rows: int = 600) -> pd.DataFrame:
    """
    차트 생성 테스트용 더미 OHLCV + 지표 컬럼 생성.
    (실데이터가 없을 때도 test_img.py가 바로 돌아가게 하기 위함)
    """
    rng = np.random.default_rng(42)
    base = 100_000 + np.cumsum(rng.normal(0, 200, size=rows))
    high = base + rng.uniform(50, 400, size=rows)
    low = base - rng.uniform(50, 400, size=rows)
    open_ = base + rng.normal(0, 50, size=rows)
    close = base + rng.normal(0, 50, size=rows)
    volume = rng.uniform(10, 1000, size=rows)

    df = pd.DataFrame(
        {
            "open": open_,
            "high": high,
            "low": low,
            "close": close,
            "volume": volume,
        }
    )

    # 간단 지표(필요 컬럼) 구성
    df["ma20"] = df["close"].rolling(20).mean()
    df["ma60"] = df["close"].rolling(60).mean()
    df["ma120"] = df["close"].rolling(120).mean()

    mid = df["close"].rolling(40).mean()
    std = df["close"].rolling(40).std()
    df["middle"] = mid
    df["upper"] = mid + (std * 2.3)
    df["lower"] = mid - (std * 2.3)

    # ma20_deg 근사
    df["ma20_deg"] = (df["ma20"] - df["ma20"].shift()).fillna(0)

    # PB, MFI10 더미
    df["PB"] = ((df["close"] - df["lower"]) / (df["upper"] - df["lower"])).replace(
        [np.inf, -np.inf], np.nan
    )
    df["PB"] = df["PB"].fillna(0.5).clip(0, 1)
    df["MFI10"] = (50 + rng.normal(0, 10, size=rows)).clip(0, 100)

    # 시간 컬럼(옵션)
    start = datetime.now() - timedelta(minutes=rows * 3)
    df["time"] = [start + timedelta(minutes=3 * i) for i in range(rows)]

    return df.ffill().fillna(0)


def _load_df(path: str | None) -> pd.DataFrame:
    if not path:
        return _dummy_df()

    if not os.path.exists(path):
        raise FileNotFoundError(path)

    if path.lower().endswith(".csv"):
        return pd.read_csv(path)
    if path.lower().endswith(".parquet"):
        return pd.read_parquet(path)

    raise ValueError("지원 포맷: .csv 또는 .parquet")


def _load_df_from_redis(
    *,
    host: str = "localhost",
    port: int = 6379,
    db: int = 0,
    key: str = "cul_df",
) -> pd.DataFrame:
    """
    main_ri_findsell_dqn_v8_real.py / RealTimeMarket와 동일하게
    Redis의 JSON(records) 형태 DataFrame을 읽어옵니다.
    """
    try:
        import redis  # type: ignore
    except ImportError as e:
        raise ImportError("Redis에서 읽으려면 redis 패키지가 필요합니다. `pip install redis`") from e

    r = redis.StrictRedis(host=host, port=port, db=db)
    raw = r.get(key)
    if not raw:
        raise KeyError(f"Redis key '{key}'에 데이터가 없습니다.")

    df_str = raw.decode("utf-8") if isinstance(raw, (bytes, bytearray)) else str(raw)
    df = pd.read_json(StringIO(df_str), orient="records")
    if df is None or df.empty:
        raise ValueError(f"Redis key '{key}'의 DataFrame이 비어있습니다.")
    return df


def _png_bytes_to_jpg_file(png_bytes: bytes, out_path: str, quality: int = 92) -> None:
    try:
        from PIL import Image
    except ImportError as e:
        raise ImportError("jpg 저장에는 Pillow가 필요합니다. `pip install pillow`") from e

    from io import BytesIO

    im = Image.open(BytesIO(png_bytes)).convert("RGB")
    im.save(out_path, format="JPEG", quality=quality, optimize=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", default=None, help="입력 데이터(.csv/.parquet). 없으면 더미 데이터 사용")
    ap.add_argument("--redis", action="store_true", help="Redis에서 cul_df를 읽어 이미지 생성")
    ap.add_argument("--redis_host", default="localhost", help="Redis host")
    ap.add_argument("--redis_port", type=int, default=6379, help="Redis port")
    ap.add_argument("--redis_db", type=int, default=0, help="Redis db index")
    ap.add_argument("--redis_key", default="cul_df", help="Redis key (기본: cul_df)")
    ap.add_argument("--outdir", default="out_imgs", help="이미지 저장 폴더")
    ap.add_argument("--ticker", default="KRW-ETH", help="차트 타이틀용 티커 문자열")
    ap.add_argument("--idx", type=int, default=500, help="df.head(idx) 기준 idx")
    ap.add_argument("--max_data", type=int, default=300, help="화면에 그릴 최대 데이터 수")
    ap.add_argument("--n_slices", type=int, default=3, help="가로 분할 조각 수")
    ap.add_argument("--overlap_px", type=int, default=0, help="조각 간 겹침 픽셀")
    ap.add_argument("--scale", type=int, default=2, help="plotly 렌더링 scale")
    ap.add_argument("--debug_env", action="store_true", help="실행 python/plotly/kaleido 환경 출력")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    if args.debug_env:
        print("python:", sys.executable)
        print("version:", sys.version.replace("\n", " "))
        try:
            import plotly

            print("plotly:", plotly.__version__, plotly.__file__)
        except Exception as e:
            print("plotly import FAIL:", repr(e))
        try:
            import kaleido

            print("kaleido:", getattr(kaleido, "__version__", None), kaleido.__file__)
        except Exception as e:
            print("kaleido import FAIL:", repr(e))

    if args.redis:
        df = _load_df_from_redis(
            host=args.redis_host,
            port=args.redis_port,
            db=args.redis_db,
            key=args.redis_key,
        )
    else:
        df = _load_df(args.csv)

    gen = ChartImageGenerator(
        ticker=args.ticker,
        config=ChartImageConfig(show=False),
    )

    base_png, parts_png = gen.render_and_slice_by_x(
        df,
        idx=args.idx,
        max_data=args.max_data,
        ticker=args.ticker,
        n_slices=args.n_slices,
        overlap_px=args.overlap_px,
        scale=args.scale,
    )

    origin_path = os.path.join(args.outdir, "origin.jpg")
    _png_bytes_to_jpg_file(base_png, origin_path)

    for i, p in enumerate(parts_png):
        out_path = os.path.join(args.outdir, f"slice_{i:02d}.jpg")
        _png_bytes_to_jpg_file(p, out_path)

    print(f"saved: {origin_path}")
    print(f"saved slices: {len(parts_png)} (dir: {args.outdir})")


if __name__ == "__main__":
    main()