result.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Any
  15. import io
  16. import pandas as pd
  17. import matplotlib.pyplot as plt
  18. from PIL import Image
  19. from ...common.result import BaseTSResult
  20. def visualize(forecast: pd.DataFrame, actual_data: pd.DataFrame) -> Image.Image:
  21. """
  22. Visualizes both the time series forecast and actual results, returning them as a Pillow image.
  23. Args:
  24. forecast (pd.DataFrame): The DataFrame containing the forecast data.
  25. actual_data (pd.Series): The actual observed data for comparison.
  26. title (str): The title of the plot.
  27. Returns:
  28. Image.Image: The visualized result as a Pillow image.
  29. """
  30. plt.figure(figsize=(12, 6))
  31. forecast_columns = forecast.columns
  32. index_name = forecast.index.name
  33. actual_data = actual_data.set_index(index_name)
  34. actual_data.index = actual_data.index.astype(str)
  35. forecast.index = forecast.index.astype(str)
  36. length = min(len(forecast), len(actual_data))
  37. actual_data = actual_data.tail(length)
  38. plt.plot(actual_data.index, actual_data[forecast_columns[0]], label='Actual Data', color='blue', linestyle='--')
  39. plt.plot(forecast.index, forecast[forecast_columns[0]], label='Forecast', color='red')
  40. plt.title('Time Series Forecast')
  41. plt.xlabel('Time')
  42. plt.ylabel(forecast_columns[0])
  43. plt.legend()
  44. plt.grid(True)
  45. plt.xticks(ticks=range(0, 2*length, 10))
  46. plt.xticks(rotation=45)
  47. buf = io.BytesIO()
  48. plt.savefig(buf, bbox_inches='tight')
  49. buf.seek(0)
  50. plt.close()
  51. image = Image.open(buf)
  52. return image
  53. class TSFcResult(BaseTSResult):
  54. """A class representing the result of a time series forecasting task."""
  55. def _to_img(self) -> Image.Image:
  56. """apply"""
  57. forecast = self["forecast"]
  58. ts_input = pd.read_csv(self["input_path"])
  59. return {"res": visualize(forecast, ts_input)}
  60. def _to_csv(self) -> Any:
  61. """
  62. Converts the forecasting results to a CSV format.
  63. Returns:
  64. Any: The forecast data formatted for CSV output, typically a DataFrame or similar structure.
  65. """
  66. return {"res": self["forecast"]}