Update app.py

app.py

@@ -7,13 +7,7 @@ import requests
 import json
 import re
 
-# 将图像转换为 base64，以便在 gradio 中显示
-def get_image_data(fig):
-    img_byte_arr = io.BytesIO() 
-    img.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
-    img_str = base64.b64encode(img_byte_arr).decode('utf-8')
-    return img_str
+
 
 # 执行 Python 代码并生成图像
 def execute_code(code):
@@ -21,7 +15,14 @@ def execute_code(code):
     exec(code, namespace)
     fig = namespace.get('fig')  # Assume the code generates a matplotlib figure named 'fig'
     if fig:
-        return get_image_data(fig)
+        img = get_image_data(fig)
+        
+        img_byte_arr = io.BytesIO()
+        img.save(img_byte_arr, format='PNG')
+        img_byte_arr = img_byte_arr.getvalue()
+        img_b64 = base64.b64encode(img_byte_arr).decode('utf-8')
+        
+        return img_b64
     else:
         raise ValueError("The code did not generate a matplotlib figure named 'fig'")