Spaces:
Running
Running
| import base64 | |
| from io import BytesIO | |
| from PIL import Image | |
| def decode_chart(charts: dict): | |
| """ | |
| Decodes base64 string charts into PIL Image objects. | |
| Also converts Matplotlib Figures into PIL Image objects. | |
| If already decoded or None, returns as is. | |
| """ | |
| if not isinstance(charts, dict): | |
| return {} | |
| result = {} | |
| for name, img_str in charts.items(): | |
| # Case 1: Base64 String (from DB) | |
| if isinstance(img_str, str) and not img_str.startswith("<matplotlib"): | |
| try: | |
| # Remove data URI prefix if present | |
| if "," in img_str: | |
| img_str = img_str.split(",")[1] | |
| img_bytes = base64.b64decode(img_str) | |
| img = Image.open(BytesIO(img_bytes)) | |
| result[name] = img | |
| except Exception as e: | |
| import logging | |
| logging.getLogger(__name__).error(f"Failed to decode base64 chart {name}: {e}") | |
| result[name] = img_str | |
| # Case 2: Matplotlib Figure (from fresh run) | |
| elif hasattr(img_str, "savefig"): | |
| try: | |
| buf = BytesIO() | |
| img_str.savefig(buf, format="png", bbox_inches='tight') | |
| buf.seek(0) | |
| img = Image.open(buf) | |
| result[name] = img | |
| except Exception as e: | |
| import logging | |
| logging.getLogger(__name__).error(f"Failed to convert figure to image for {name}: {e}") | |
| result[name] = img_str | |
| else: | |
| result[name] = img_str | |
| return result | |
| def figure_to_base64(fig): | |
| """ | |
| Converts a Matplotlib Figure to a base64 string for database storage. | |
| """ | |
| buf = BytesIO() | |
| fig.savefig(buf, format='png', bbox_inches='tight') | |
| buf.seek(0) | |
| img_str = base64.b64encode(buf.getvalue()).decode('utf-8') | |
| return img_str | |
| def serialize_charts(charts: dict) -> dict: | |
| """ | |
| Ensures all charts in the dict are JSON-serializable base64 strings. | |
| Handles: | |
| - str: Assumed to be already base64. | |
| - plt.Figure: Serialized using savefig. | |
| - PIL.Image: Serialized using img.save. | |
| """ | |
| if not isinstance(charts, dict): | |
| return {} | |
| result = {} | |
| for name, val in charts.items(): | |
| if val is None: | |
| continue | |
| if isinstance(val, str): | |
| # already serialized | |
| result[name] = val | |
| elif hasattr(val, "savefig"): | |
| # Matplotlib Figure | |
| result[name] = figure_to_base64(val) | |
| elif hasattr(val, "save"): | |
| # PIL Image | |
| try: | |
| buf = BytesIO() | |
| val.save(buf, format="png") | |
| buf.seek(0) | |
| result[name] = base64.b64encode(buf.getvalue()).decode("utf-8") | |
| except Exception as e: | |
| import logging | |
| logging.getLogger(__name__).error(f"Failed to serialize PIL image for {name}: {e}") | |
| result[name] = None | |
| else: | |
| result[name] = str(val) | |
| return result |