team-10/image_generation.py
2025-08-02 12:26:50 +02:00

58 lines
1.9 KiB
Python

import os
import pandas as pd
import random
from streamlit import echo
METADATA_CSV = ".\plant_growth_data\plant_growth_data.csv"
df = pd.read_csv(METADATA_CSV)
DATASET_DIR = ".\plant_growth_data\dataset"
PLANTVILLAGE_DIR = ".\plant_growth_data\images"
def determine_health_status(description: str) -> str:
illness_keywords = [
"disease", "wilt", "wilted ", "brown", "yellowing", "withered", "damage", "ill", "unhealthy", "harmful", "fire", "coca cola"
]
healthy_keywords = [
"healthy", "vibrant", "green", "lush", "strong", "thriving", "fresh", "good condition"
]
desc_lower = description.lower()
if any(word in desc_lower for word in healthy_keywords):
return "healthy"
if any(word in desc_lower for word in illness_keywords):
return "ill"
return "healthy"
def filter_images_for_row(plant_folder):
images = os.listdir(plant_folder)
return images
def generate_image(plant, description, plant_age):
plant_health_status = determine_health_status(description)
plant_map = {
"Basil": "basil\\" + plant_health_status,
"Tomato": "tomato\\" + plant_health_status,
"Lettuce": "lettuce\\" + plant_health_status,
}
folder_name = plant_map.get(plant)
if folder_name is None:
raise ValueError(f"Pianta '{plant}' non trovata nel dataset.")
plant_folder = os.path.join(PLANTVILLAGE_DIR, folder_name)
if not os.path.exists(plant_folder):
raise FileNotFoundError(f"Cartella {plant_folder} non trovata.")
candidate_images = filter_images_for_row(plant_folder)
if not candidate_images:
raise FileNotFoundError(f"Nessuna immagine trovata nella cartella {plant_folder}.")
if plant_age < 30:
chosen_image = [img for img in candidate_images if "young" in img]
else:
chosen_image = [img for img in candidate_images if "mature" in img]
image_path = os.path.join(plant_folder, chosen_image[0])
return image_path