Yoloを使って株価チャートの予測(?)をしてみる

元ネタは、2023年にアップロードされたstockmark-future-prediction

foduucom/stockmarket-future-prediction · Hugging Face

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

というものです。

hugging faceを漁っていたら見つけました。試したこと無かったので、試してみます。

内容としては、株価チャートが与えられたら、その後のトレンドを検出するモデルのようです。

foduucom/product-detection-in-shelf-yolov8 — https://huggingface.co/foduucom/stockmarket-future-predictionより

他に検証している人が見当たらなかったので、自分で動かしてみることにしました。

実行環境

Dockerfile

前に作ったhugging faceの実行環境をベースに構築しました。yoloしか入れないので、ここまでいろいろと入れなくても、動くと思います。

FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04

ARG USERNAME=vscode
ARG USER_UID=1000
ARG USER_GID=$USER_UID

ENV LANG ja_JP.UTF-8
ENV LANGUAGE ja_JP:ja
ENV LC_ALL ja_JP.UTF-8
ENV TZ JST-9
ENV TERM xterm

RUN apt-get update \
    && groupadd --gid $USER_GID $USERNAME \
    && useradd -s /bin/bash --uid $USER_UID --gid $USER_GID -m $USERNAME \
    && apt-get install -y sudo \
    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
    && chmod 0440 /etc/sudoers.d/$USERNAME \
    && apt-get -y install locales \
    && localedef -f UTF-8 -i ja_JP ja_JP.UTF-8

RUN apt-get -y update \
    && apt-get install -y software-properties-common \
    && add-apt-repository ppa:deadsnakes/ppa

RUN apt install -y bash \
    build-essential \
    git \
    git-lfs \
    curl \
    ca-certificates \
    libsndfile1-dev \
    libgl1 \
    python3.10 \
    python3-pip \
    python3.10-venv && \
    rm -rf /var/lib/apt/lists

# make sure to use venv
RUN python3.10 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \
    python3.10 -m uv pip install --no-cache-dir \
    torch \
    torchvision \
    torchaudio \
    invisible_watermark && \
    python3.10 -m pip install --no-cache-dir \
    accelerate \
    datasets \
    hf-doc-builder \
    huggingface-hub \
    Jinja2 \
    librosa \
    numpy \
    scipy \
    tensorboard \
    transformers \
    pytorch-lightning

RUN pip install ultralyticsplus==0.0.28 ultralytics==8.0.43

devcontainer.json

名前とサービス名を決定してください。また、必要な拡張機能を入れてください。筆者は、pythonの拡張機能とjupyter notebookの拡張機能をいつも入れています。

{
    "name": "yolov8s_stockmarket_future_prediction-test",
    "service": "yolov8s_stockmarket_future_prediction-test",
    "dockerComposeFile": "docker-compose.yml",
    "remoteUser": "vscode",
    "workspaceFolder": "/work",
    "customizations": {
      "vscode": {
        "extensions": [
          "ms-python.python",
          "ms-toolsai.jupyter"
        ]
      }
    }
}

docker-compose.yml

設定したサービス名を間違えないでください。

また、GPUを使うので、GPUが使えるようにdeployを書き足します。

version: '3'
services:
  yolov8s_stockmarket_future_prediction-test:
    container_name: 'yolov8s_stockmarket_future_prediction-test-container'
    hostname: 'yolov8s_stockmarket_future_prediction-test-container'
    build: .
    restart: always
    working_dir: '/work' 
    tty: true
    volumes:
      - type: bind
        source: ..
        target: /work
    ulimits:
      memlock: -1
      stack: -1
    shm_size: '10gb'
    deploy:
      resources:
          reservations:
              devices:
                - capabilities: [gpu]

実行プログラム

元のURLにもHow to Get Startが書かれているので、それを参考にテストコードを書いてみます。

リモートDocker環境なので、renderingはしません。predictの引数に保存フラグを立てておけば、検出画像が出力されます。

from ultralyticsplus import YOLO, render_result
import cv2

# load model
model = YOLO('foduucom/stockmarket-future-prediction')

# set model parameters
model.overrides['conf'] = 0.25  # NMS confidence threshold
model.overrides['iou'] = 0.45  # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic
model.overrides['max_det'] = 1000  # maximum number of detections per image

# set image
image = 'images/week_chart_02.png'

# perform inference
results = model.predict(image, save=True, save_txt=True, save_conf=True)

# observe results
print(results[0].boxes)

特に、エラーもなく実行できました。モデルもhugging_faceから自動でダウンロードしてくれるようです。

次に、チャート画像を用意する必要があります。

YAHOOファイナンスから、日経平均の株価チャートをスクショしてきました。次の五枚です。