avfranco's picture
HF Space deploy snapshot (minimal allow-list)
d64fd55
Raw
History Blame Contribute Delete
3.09 kB
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