Spaces:
Running
Running
File size: 6,487 Bytes
572d3da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
"""
Visualizer Agent for PaperBanana framework.
Renders academic illustrations using image generation models.
Supports both diagram generation and statistical plot generation.
"""
import os
import mimetypes
from typing import Optional
from google import genai
from google.genai import types
import config
from utils import save_binary_file
class VisualizerAgent:
"""
Visualizer Agent: Renders illustrations from textual descriptions.
Supports two modes:
1. Diagram mode: Uses image generation model (Nano-Banana-Pro / Gemini Image)
2. Plot mode: Generates Python Matplotlib code for statistical plots
"""
def __init__(self, mode: str = "diagram"):
"""
Initialize Visualizer Agent.
Args:
mode: Generation mode - "diagram" or "plot"
"""
self.client = genai.Client(api_key=config.GEMINI_API_KEY)
self.mode = mode
if mode == "diagram":
self.model = config.IMAGE_MODEL
elif mode == "plot":
self.model = config.VLM_MODEL # Use VLM for code generation
else:
raise ValueError(f"Invalid mode: {mode}. Use 'diagram' or 'plot'")
def visualize(self,
description: str,
output_path: str = "output",
data: dict = None) -> str:
"""
Generate visualization from description.
Args:
description: Textual description of the illustration
output_path: Base path for output file (without extension)
data: Optional data dict for plot mode
Returns:
Path to generated image file or code file
"""
if self.mode == "diagram":
return self._generate_diagram(description, output_path)
elif self.mode == "plot":
return self._generate_plot(description, output_path, data)
def _generate_diagram(self, description: str, output_path: str) -> str:
"""
Generate diagram image using image generation model.
Args:
description: Detailed visual description
output_path: Base path for output file
Returns:
Path to generated image
"""
# Create prompt for image generation
prompt = f"""Generate a high-quality academic methodology diagram with the following specifications:
{description}
Requirements:
- Professional academic publication quality
- Clear, readable text and labels
- Consistent styling throughout
- Appropriate use of colors and shapes
- Publication-ready resolution
"""
contents = [
types.Content(
role="user",
parts=[types.Part.from_text(text=prompt)]
)
]
generate_config = types.GenerateContentConfig(
response_modalities=["IMAGE", "TEXT"],
image_config=types.ImageConfig(
image_size=config.IMAGE_SIZE
)
)
file_index = 0
saved_path = None
for chunk in self.client.models.generate_content_stream(
model=self.model,
contents=contents,
config=generate_config
):
if (chunk.candidates is None or
chunk.candidates[0].content is None or
chunk.candidates[0].content.parts is None):
continue
# Check for inline image data
part = chunk.candidates[0].content.parts[0]
if part.inline_data and part.inline_data.data:
inline_data = part.inline_data
data_buffer = inline_data.data
file_extension = mimetypes.guess_extension(inline_data.mime_type)
if file_extension:
file_name = f"{output_path}_{file_index}{file_extension}"
saved_path = save_binary_file(file_name, data_buffer)
file_index += 1
else:
# Print any text output
if chunk.text:
print(chunk.text)
return saved_path or f"{output_path}_0.png"
def _generate_plot(self, description: str, output_path: str, data: dict = None) -> str:
"""
Generate statistical plot by creating Matplotlib code.
Args:
description: Description of desired plot
output_path: Base path for output code file
data: Optional data dictionary
Returns:
Path to generated Python code file
"""
data_context = ""
if data:
data_context = f"\n\nDATA PROVIDED:\n{str(data)}\n"
prompt = f"""You are an expert at creating publication-quality statistical plots using Matplotlib.
Generate complete, executable Python code using Matplotlib to create the following plot:
{description}
{data_context}
Requirements:
1. Use professional academic styling (seaborn-paper style or similar)
2. Include clear axis labels with units
3. Add legend if multiple series
4. Use appropriate colors and markers
5. Set figure size for publication (e.g., 6x4 inches)
6. Save as high-resolution PNG (300 dpi minimum)
7. Include error bars if applicable
8. Follow best practices for data visualization
OUTPUT FORMAT:
Provide ONLY the complete Python code, ready to execute.
Start with necessary imports and end with plt.savefig().
Do not include any explanations outside the code comments.
"""
contents = [
types.Content(
role="user",
parts=[types.Part.from_text(text=prompt)]
)
]
generate_config = types.GenerateContentConfig(
thinking_config=types.ThinkingConfig(
thinking_level="MEDIUM"
)
)
code = ""
for chunk in self.client.models.generate_content_stream(
model=self.model,
contents=contents,
config=generate_config
):
code += chunk.text
# Save code to file
code_file = f"{output_path}.py"
with open(code_file, 'w') as f:
f.write(code.strip())
print(f"Plot code saved to: {code_file}")
print("Run the code to generate the plot image.")
return code_file
|