In [ ]:
# 生成K线图并保存为图像
def generate_chart_image(data, window_size=50, output_path='chart.png'):
plt.figure(figsize=(10, 6))
# 绘制K线图
for i in range(len(data)-window_size, len(data)):
date = data.index[i]
op, hi, lo, cl = data.loc[date, ['Open', 'High', 'Low', 'Close']]
# 绿色表示收盘价高于开盘价，红色表示开盘价高于收盘价
color = 'g' if cl >= op else 'r'
plt.plot([i, i], [lo, hi], color=color)
plt.plot([i, i-0.2], [op, op], color=color)
plt.plot([i, i+0.2], [cl, cl], color=color)
plt.title('股价图表')
plt.xticks([]) # 移除x轴标签以简化图像
plt.savefig(output_path)
plt.close()
# 读取并预处理图像
image = cv2.imread(output_path)
image = cv2.resize(image, (224, 224)) # 调整为模型输入尺寸
return image
# 生成训练数据
windows = []
labels = [] # 假设有预定义的形态标签
for i in range(len(stock_data) - 50):
window_data = stock_data.iloc[i:i+50]
image = generate_chart_image(window_data)
windows.append(image)
# 这里需要实际的标签数据，可以是手动标注或使用规则自动生成
# labels.append(...)
# 通过未来5个交易日的价格变动生成标签
if i + 55 < len(stock_data):
future_return = (stock_data.iloc[i+55]['Close'] / stock_data.iloc[i+50]['Close'] - 1) * 100
if future_return > 3:
pattern_label = 0 # 大幅上涨
elif future_return > 1:
pattern_label = 1 # 小幅上涨
elif future_return < -3:
pattern_label = 2 # 大幅下跌
elif future_return < -1:
pattern_label = 3 # 小幅下跌
else:
pattern_label = 4 # 盘整
labels.append(pattern_label)
else:
# 对于末尾没有足够未来数据的窗口，移除对应的X数据
windows.pop()
X = np.array(windows)
y = np.array(labels)
# 可视化一个样本图表（可选）
plt.figure(figsize=(8, 8))
plt.imshow(cv2.cvtColor(X[0], cv2.COLOR_BGR2RGB))
plt.title('比亚迪K线图样本')
plt.axis('off')
plt.show()
