240 lines
9.8 KiB
Python
240 lines
9.8 KiB
Python
|
from PIL import Image, ImageDraw, ImageFont
|
||
|
|
import random
|
||
|
|
import math
|
||
|
|
import os
|
||
|
|
|
||
|
|
class ImageGenerator:
|
||
|
|
def __init__(self):
|
||
|
|
self.image_size = (400, 300)
|
||
|
|
self.plant_colors = {
|
||
|
|
'tomato': {'stem': '#228B22', 'leaf': '#32CD32', 'fruit': '#FF6347'},
|
||
|
|
'basil': {'stem': '#228B22', 'leaf': '#90EE90', 'fruit': '#FFFFFF'},
|
||
|
|
'mint': {'stem': '#228B22', 'leaf': '#98FB98', 'fruit': '#FFFFFF'},
|
||
|
|
'lettuce': {'stem': '#228B22', 'leaf': '#ADFF2F', 'fruit': '#FFFFFF'},
|
||
|
|
'rosemary': {'stem': '#8B4513', 'leaf': '#556B2F', 'fruit': '#FFFFFF'},
|
||
|
|
'strawberry': {'stem': '#228B22', 'leaf': '#32CD32', 'fruit': '#FF1493'}
|
||
|
|
}
|
||
|
|
|
||
|
|
def generate_evolution(self, plant_type, parameters, prediction, baseline_path=None):
|
||
|
|
"""Generate a series of images showing plant evolution"""
|
||
|
|
growth_stages = prediction['growth_stages']
|
||
|
|
health_score = prediction['health_score']
|
||
|
|
|
||
|
|
images = []
|
||
|
|
|
||
|
|
# Generate images for key growth stages
|
||
|
|
stage_indices = [0, len(growth_stages)//4, len(growth_stages)//2,
|
||
|
|
3*len(growth_stages)//4, len(growth_stages)-1]
|
||
|
|
|
||
|
|
for i, stage_idx in enumerate(stage_indices):
|
||
|
|
if stage_idx < len(growth_stages):
|
||
|
|
height = growth_stages[stage_idx]
|
||
|
|
image = self.generate_plant_image(plant_type, height, health_score, i+1)
|
||
|
|
images.append(image)
|
||
|
|
|
||
|
|
return images
|
||
|
|
|
||
|
|
def generate_plant_image(self, plant_type, height, health_score, stage):
|
||
|
|
"""Generate a single plant image"""
|
||
|
|
img = Image.new('RGB', self.image_size, color='#87CEEB') # Sky blue background
|
||
|
|
draw = ImageDraw.Draw(img)
|
||
|
|
|
||
|
|
# Draw ground
|
||
|
|
ground_y = self.image_size[1] - 50
|
||
|
|
draw.rectangle([0, ground_y, self.image_size[0], self.image_size[1]],
|
||
|
|
fill='#8B4513') # Brown ground
|
||
|
|
|
||
|
|
# Get plant colors (default to tomato since we removed plant type selection)
|
||
|
|
colors = self.plant_colors.get(plant_type, self.plant_colors['tomato'])
|
||
|
|
|
||
|
|
# Calculate plant dimensions based on height and health
|
||
|
|
if stage == 0: # Initial/seed stage
|
||
|
|
plant_height = 5 # Very small initial plant
|
||
|
|
plant_width = 3
|
||
|
|
else:
|
||
|
|
plant_height = min(height * 2, self.image_size[1] - 100) # Scale for display
|
||
|
|
plant_width = plant_height * 0.6
|
||
|
|
|
||
|
|
# Adjust colors based on health
|
||
|
|
health_factor = health_score / 100.0
|
||
|
|
stem_color = self._adjust_color_health(colors['stem'], health_factor)
|
||
|
|
leaf_color = self._adjust_color_health(colors['leaf'], health_factor)
|
||
|
|
|
||
|
|
# Plant center position
|
||
|
|
center_x = self.image_size[0] // 2
|
||
|
|
base_y = ground_y
|
||
|
|
|
||
|
|
if stage == 0:
|
||
|
|
# Draw seed/initial state
|
||
|
|
self._draw_seed(draw, center_x, base_y - 10, health_factor)
|
||
|
|
else:
|
||
|
|
# Draw stem
|
||
|
|
stem_width = max(2, int(plant_height * 0.05))
|
||
|
|
stem_top_y = base_y - plant_height
|
||
|
|
draw.rectangle([center_x - stem_width//2, int(stem_top_y),
|
||
|
|
center_x + stem_width//2, base_y], fill=stem_color)
|
||
|
|
|
||
|
|
# Draw leaves based on plant type and stage
|
||
|
|
self._draw_leaves(draw, plant_type, center_x, stem_top_y, base_y,
|
||
|
|
plant_width, leaf_color, stage)
|
||
|
|
|
||
|
|
# Draw fruits/flowers if applicable
|
||
|
|
if stage >= 3 and plant_type in ['tomato', 'strawberry']:
|
||
|
|
self._draw_fruits(draw, plant_type, center_x, stem_top_y, base_y,
|
||
|
|
colors['fruit'], stage)
|
||
|
|
|
||
|
|
# Add stage label
|
||
|
|
try:
|
||
|
|
font = ImageFont.load_default()
|
||
|
|
except:
|
||
|
|
font = None
|
||
|
|
|
||
|
|
if stage == 0:
|
||
|
|
stage_text = f"Initial State - Seed"
|
||
|
|
else:
|
||
|
|
stage_text = f"Stage {stage} - Height: {height:.1f}cm"
|
||
|
|
|
||
|
|
if font:
|
||
|
|
draw.text((10, 10), stage_text, fill='black', font=font)
|
||
|
|
else:
|
||
|
|
draw.text((10, 10), stage_text, fill='black')
|
||
|
|
|
||
|
|
# Add health indicator
|
||
|
|
health_text = f"Health: {health_score:.1f}%"
|
||
|
|
health_color = 'green' if health_score > 70 else 'orange' if health_score > 40 else 'red'
|
||
|
|
if font:
|
||
|
|
draw.text((10, 30), health_text, fill=health_color, font=font)
|
||
|
|
else:
|
||
|
|
draw.text((10, 30), health_text, fill=health_color)
|
||
|
|
|
||
|
|
return img
|
||
|
|
|
||
|
|
def _draw_seed(self, draw, x, y, health_factor):
|
||
|
|
"""Draw a seed for the initial state"""
|
||
|
|
seed_size = 8
|
||
|
|
seed_color = '#8B4513' # Brown seed
|
||
|
|
|
||
|
|
# Adjust seed color based on health
|
||
|
|
if health_factor > 0.7:
|
||
|
|
seed_color = '#654321' # Healthy brown
|
||
|
|
elif health_factor > 0.4:
|
||
|
|
seed_color = '#8B4513' # Normal brown
|
||
|
|
else:
|
||
|
|
seed_color = '#A0522D' # Pale brown
|
||
|
|
|
||
|
|
# Draw seed
|
||
|
|
draw.ellipse([x - seed_size, y - seed_size//2,
|
||
|
|
x + seed_size, y + seed_size//2], fill=seed_color)
|
||
|
|
|
||
|
|
# Draw small sprout if health is good
|
||
|
|
if health_factor > 0.5:
|
||
|
|
sprout_color = '#90EE90'
|
||
|
|
draw.line([x, y - seed_size//2, x, y - seed_size//2 - 5],
|
||
|
|
fill=sprout_color, width=2)
|
||
|
|
|
||
|
|
def _draw_leaves(self, draw, plant_type, center_x, top_y, base_y, width, color, stage):
|
||
|
|
"""Draw leaves based on plant type"""
|
||
|
|
plant_height = base_y - top_y
|
||
|
|
num_leaves = min(stage * 2, 8) # More leaves as plant grows
|
||
|
|
|
||
|
|
for i in range(num_leaves):
|
||
|
|
# Calculate leaf position
|
||
|
|
y_pos = base_y - (i + 1) * (plant_height / (num_leaves + 1))
|
||
|
|
side = 1 if i % 2 == 0 else -1 # Alternate sides
|
||
|
|
|
||
|
|
leaf_x = center_x + side * (width * 0.3)
|
||
|
|
leaf_size = width * 0.2 * (1 + stage * 0.1)
|
||
|
|
|
||
|
|
if plant_type == 'lettuce':
|
||
|
|
# Draw broad leaves for lettuce
|
||
|
|
self._draw_broad_leaf(draw, leaf_x, y_pos, leaf_size, color)
|
||
|
|
elif plant_type == 'rosemary':
|
||
|
|
# Draw needle-like leaves for rosemary
|
||
|
|
self._draw_needle_leaf(draw, leaf_x, y_pos, leaf_size, color)
|
||
|
|
else:
|
||
|
|
# Draw regular oval leaves (default for tomato)
|
||
|
|
self._draw_oval_leaf(draw, leaf_x, y_pos, leaf_size, color)
|
||
|
|
|
||
|
|
def _draw_oval_leaf(self, draw, x, y, size, color):
|
||
|
|
"""Draw an oval leaf"""
|
||
|
|
draw.ellipse([x - size//2, y - size//3, x + size//2, y + size//3], fill=color)
|
||
|
|
|
||
|
|
def _draw_broad_leaf(self, draw, x, y, size, color):
|
||
|
|
"""Draw a broad leaf for lettuce"""
|
||
|
|
points = [
|
||
|
|
(x, y - size//2),
|
||
|
|
(x + size//2, y),
|
||
|
|
(x, y + size//2),
|
||
|
|
(x - size//2, y)
|
||
|
|
]
|
||
|
|
draw.polygon(points, fill=color)
|
||
|
|
|
||
|
|
def _draw_needle_leaf(self, draw, x, y, size, color):
|
||
|
|
"""Draw needle-like leaves for rosemary"""
|
||
|
|
for i in range(3):
|
||
|
|
offset = (i - 1) * 3
|
||
|
|
draw.line([x + offset, y - size//4, x + offset, y + size//4],
|
||
|
|
fill=color, width=2)
|
||
|
|
|
||
|
|
def _draw_fruits(self, draw, plant_type, center_x, top_y, base_y, color, stage):
|
||
|
|
"""Draw fruits based on plant type"""
|
||
|
|
if plant_type == 'tomato':
|
||
|
|
# Draw tomatoes
|
||
|
|
num_fruits = min(stage - 2, 4)
|
||
|
|
for i in range(num_fruits):
|
||
|
|
fruit_x = center_x + random.randint(-20, 20)
|
||
|
|
fruit_y = int(top_y + random.randint(10, (base_y - top_y) // 2))
|
||
|
|
fruit_size = 8 + stage * 2
|
||
|
|
draw.ellipse([fruit_x - fruit_size, fruit_y - fruit_size,
|
||
|
|
fruit_x + fruit_size, fruit_y + fruit_size], fill=color)
|
||
|
|
|
||
|
|
elif plant_type == 'strawberry':
|
||
|
|
# Draw strawberries
|
||
|
|
num_fruits = min(stage - 2, 3)
|
||
|
|
for i in range(num_fruits):
|
||
|
|
fruit_x = center_x + random.randint(-15, 15)
|
||
|
|
fruit_y = base_y - random.randint(20, 40)
|
||
|
|
self._draw_strawberry(draw, fruit_x, fruit_y, color)
|
||
|
|
|
||
|
|
def _draw_strawberry(self, draw, x, y, color):
|
||
|
|
"""Draw a strawberry shape"""
|
||
|
|
# Draw strawberry body
|
||
|
|
points = [(x, y - 8), (x + 6, y), (x, y + 8), (x - 6, y)]
|
||
|
|
draw.polygon(points, fill=color)
|
||
|
|
|
||
|
|
# Draw strawberry top (green)
|
||
|
|
draw.polygon([(x - 3, y - 8), (x, y - 12), (x + 3, y - 8)], fill='green')
|
||
|
|
|
||
|
|
def _adjust_color_health(self, color_hex, health_factor):
|
||
|
|
"""Adjust color based on plant health"""
|
||
|
|
# Convert hex to RGB
|
||
|
|
color_hex = color_hex.lstrip('#')
|
||
|
|
r, g, b = tuple(int(color_hex[i:i+2], 16) for i in (0, 2, 4))
|
||
|
|
|
||
|
|
# Adjust brightness based on health
|
||
|
|
factor = 0.5 + health_factor * 0.5 # Range from 0.5 to 1.0
|
||
|
|
r = int(r * factor)
|
||
|
|
g = int(g * factor)
|
||
|
|
b = int(b * factor)
|
||
|
|
|
||
|
|
# Ensure values are within valid range
|
||
|
|
r = max(0, min(255, r))
|
||
|
|
g = max(0, min(255, g))
|
||
|
|
b = max(0, min(255, b))
|
||
|
|
|
||
|
|
return f'#{r:02x}{g:02x}{b:02x}'
|
||
|
|
|
||
|
|
def save_evolution_sequence(self, images, filename):
|
||
|
|
"""Save evolution images as separate files"""
|
||
|
|
if images:
|
||
|
|
try:
|
||
|
|
base_name = filename.rsplit('.', 1)[0]
|
||
|
|
for i, image in enumerate(images):
|
||
|
|
stage_filename = f"{base_name}_stage_{i+1}.png"
|
||
|
|
image.save(stage_filename)
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error saving images: {e}")
|
||
|
|
return False
|
||
|
|
return False
|