Loading vagen/env/__init__.py +3 −2 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanEnvConfig from .frozenlake import FrozenLakeEnv,FrozenLakeEnvConfig, FrozenLakeService from .navigation import NavigationEnv, NavigationEnvConfig, NavigationServiceConfig, NavigationService from .svg import SVGEnv, SvgEnvConfig, SVGService from .svg import SVGEnv, SvgEnvConfig, SVGService, SVGServiceConfig from .primitive_skill import PrimitiveSkillEnv, PrimitiveSkillEnvConfig, PrimitiveSkillService, PrimitiveSkillConfig REGISTERED_ENV = { "sokoban": { Loading @@ -22,7 +22,8 @@ REGISTERED_ENV = { "svg": { "env_cls": SVGEnv, "config_cls": SvgEnvConfig, "service_cls": SVGService "service_cls": SVGService, "service_config_cls": SVGServiceConfig }, "primitive_skill": { "env_cls": PrimitiveSkillEnv, Loading vagen/env/svg/__init__.py +1 −0 Original line number Diff line number Diff line from .env_config import SvgEnvConfig from .env import SVGEnv from .service import SVGService from .service_config import SVGServiceConfig No newline at end of file vagen/env/svg/dreamsim.py +34 −0 Original line number Diff line number Diff line import torch from PIL import Image import os from dreamsim import dreamsim # Create global cache and lock, similar to DINO implementation _model_cache = {} _model_cache_lock = threading.Lock() _model_counter = 0 def get_dreamsim_model(device=None): """ Get a singleton instance of DreamSim model, using cache to avoid duplicate loading Args: device: Device to run model on Returns: DreamSimScoreCalculator: Instance of DreamSim calculator """ global _model_counter # Choose device based on availability if not specified if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" # Use device as cache key cache_key = f"dreamsim_{device}" with _model_cache_lock: if cache_key not in _model_cache: _model_counter += 1 pid = os.getpid() logging.info(f"Process {pid}: Created DreamSim model #{_model_counter} on {device}") _model_cache[cache_key] = DreamSimScoreCalculator(device=device) return _model_cache[cache_key] class DreamSimScoreCalculator: """ A wrapper class for DreamSim model to calculate similarity scores between images. Loading @@ -15,6 +48,7 @@ class DreamSimScoreCalculator: cache_dir: Cache directory for model weights device: Device to run the model on (defaults to CUDA if available, else CPU) """ cache_dir = os.path.expanduser(cache_dir) if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: Loading vagen/env/svg/score.py +78 −48 Original line number Diff line number Diff line import numpy as np import cv2 import os from vagen.env.svg.dino import DINOScoreCalculator Loading Loading @@ -98,82 +97,113 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m return scores def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_model=None, dreamsim_model=None): def calculate_total_score_batch(gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model=None, dreamsim_model=None): """ Calculate all metrics and return a comprehensive score Calculate scores for multiple image pairs in batch mode Args: gt_im: Ground truth image gen_im: Generated image gt_code: Ground truth SVG code gen_code: Generated SVG code score_config: Dictionary containing scoring parameters - model_size: small, base, large - dino_only: Whether to use only DINO for scoring - dino_weight: Weight for DINO score - structural_weight: Weight for structural score - dreamsim_weight: Weight for DreamSim score (new) gt_images: List of ground truth images gen_images: List of generated images gt_codes: List of ground truth SVG codes gen_codes: List of generated SVG codes score_configs: List of scoring parameters dictionaries dino_model: Pre-loaded DINO model (optional) dreamsim_model: Pre-loaded DreamSim model (optional) Returns: dict: Dictionary of all scores including the total weighted score List of dictionaries containing all scores """ # Get configuration parameters with defaults model_size = score_config.get("model_size", "small") dino_only = score_config.get("dino_only", False) # Define default weights based on model size default_weights = { "small": {"dino": 3.0, "structural": 7.0, "dreamsim": 5.0}, "base": {"dino": 5.0, "structural": 5.0, "dreamsim": 5.0}, "large": {"dino": 6.0, "structural": 4.0, "dreamsim": 5.0} } batch_size = len(gt_images) if batch_size == 0: return [] # Get weights with defaults weights = { "dino": score_config.get("dino_weight", default_weights[model_size]["dino"]), "structural": score_config.get("structural_weight", default_weights[model_size]["structural"]), "dreamsim": score_config.get("dreamsim_weight", default_weights[model_size]["dreamsim"]) } # Verify all inputs have same batch size if not (len(gen_images) == len(gt_codes) == len(gen_codes) == len(score_configs) == batch_size): raise ValueError("All input lists must have the same length") # Initialize scores scores = { # Initialize results batch_results = [{ "dino_score": 0.0, "structural_score": 0.0, "dreamsim_score": 0.0, "total_score": 0.0 } } for _ in range(batch_size)] # Calculate DINO score if needed if weights["dino"] > 0: # Check if we need to calculate DINO scores need_dino = any(score_config.get("dino_weight", 5.0) > 0 for score_config in score_configs) # Check if we need to calculate DreamSim scores need_dreamsim = any(score_config.get("dreamsim_weight", 5.0) > 0 for score_config in score_configs) # Calculate DINO scores in batch if needed if need_dino: if dino_model is None: from vagen.env.svg.dino import get_dino_model # Default to small model size if not specified model_size = score_configs[0].get("model_size", "small") if score_configs else "small" dino_model = get_dino_model(model_size) scores["dino_score"] = float(dino_model.calculate_DINOv2_similarity_score(gt_im=gt_im, gen_im=gen_im)) # Calculate DreamSim score if needed if weights["dreamsim"] > 0: # Calculate all DINO scores at once using batch processing dino_scores = dino_model.calculate_batch_scores(gt_images, gen_images) # Assign scores to results for i, score in enumerate(dino_scores): batch_results[i]["dino_score"] = float(score) # Calculate DreamSim scores in batch if needed if need_dreamsim: if dreamsim_model is None: from vagen.env.svg.dreamsim import get_dreamsim_model dreamsim_model = get_dreamsim_model() scores["dreamsim_score"] = float(dreamsim_model.calculate_similarity_score(gt_im=gt_im, gen_im=gen_im)) # If DINO only mode, return only DINO score # Calculate all DreamSim scores at once using batch processing dreamsim_scores = dreamsim_model.calculate_batch_scores(gt_images, gen_images) # Assign scores to results for i, score in enumerate(dreamsim_scores): batch_results[i]["dreamsim_score"] = float(score) # Calculate structural scores and total scores for i in range(batch_size): score_config = score_configs[i] result = batch_results[i] # Check if DINO-only mode dino_only = score_config.get("dino_only", False) if dino_only: scores["total_score"] = scores["dino_score"] return scores result["total_score"] = result["dino_score"] continue # Get model size for default weights model_size = score_config.get("model_size", "small") # Define default weights based on model size default_weights = { "small": {"dino": 3.0, "structural": 7.0, "dreamsim": 5.0}, "base": {"dino": 5.0, "structural": 5.0, "dreamsim": 5.0}, "large": {"dino": 6.0, "structural": 4.0, "dreamsim": 5.0} } # Get weights with defaults weights = { "dino": score_config.get("dino_weight", default_weights[model_size]["dino"]), "structural": score_config.get("structural_weight", default_weights[model_size]["structural"]), "dreamsim": score_config.get("dreamsim_weight", default_weights[model_size]["dreamsim"]) } # Calculate structural score if needed if weights["structural"] > 0: scores["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_im, gen_im))) from vagen.env.svg.score import calculate_structural_accuracy result["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_images[i], gen_images[i]))) # Calculate weighted total score weighted_sum = ( scores["dino_score"] * weights["dino"] + scores["structural_score"] * weights["structural"] + scores["dreamsim_score"] * weights["dreamsim"] result["dino_score"] * weights["dino"] + result["structural_score"] * weights["structural"] + result["dreamsim_score"] * weights["dreamsim"] ) scores['total_score'] = max(0.0, weighted_sum) result["total_score"] = max(0.0, weighted_sum) return scores No newline at end of file return batch_results No newline at end of file vagen/env/svg/service.py +15 −1 Original line number Diff line number Diff line Loading @@ -37,6 +37,9 @@ class SVGService(BaseService): self.model_size = self.config.model_size self.dino_model = None # Will be loaded on first use # Add DreamSim model support self.dreamsim_model = None # Will be loaded on first use if enabled # Store device for model inference self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"SVGService initialized with {self.max_workers} workers, model_size={self.model_size}, device={self.device}") Loading @@ -50,6 +53,16 @@ class SVGService(BaseService): self.dino_model = get_dino_model(self.model_size, self.device) return self.dino_model def _get_dreamsim_model(self): """ Get or initialize the DreamSim model. Uses lazy loading to avoid loading the model until needed. """ if self.dreamsim_model is None: from vagen.env.svg.dreamsim import get_dreamsim_model self.dreamsim_model = get_dreamsim_model(self.device) return self.dreamsim_model def create_environments_batch(self, ids2configs: Dict[Any, Any]) -> None: """ Create multiple SVG environments in parallel. Loading Loading @@ -181,10 +194,11 @@ class SVGService(BaseService): if valid_env_ids: # Get DINO model dino_model = self._get_dino_model() dreamsim_model = self._get_dreamsim_model() # Calculate all scores at once batch_results = calculate_total_score_batch( gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model, dreamsim_model ) # Process results directly using the index mapping Loading Loading
vagen/env/__init__.py +3 −2 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanEnvConfig from .frozenlake import FrozenLakeEnv,FrozenLakeEnvConfig, FrozenLakeService from .navigation import NavigationEnv, NavigationEnvConfig, NavigationServiceConfig, NavigationService from .svg import SVGEnv, SvgEnvConfig, SVGService from .svg import SVGEnv, SvgEnvConfig, SVGService, SVGServiceConfig from .primitive_skill import PrimitiveSkillEnv, PrimitiveSkillEnvConfig, PrimitiveSkillService, PrimitiveSkillConfig REGISTERED_ENV = { "sokoban": { Loading @@ -22,7 +22,8 @@ REGISTERED_ENV = { "svg": { "env_cls": SVGEnv, "config_cls": SvgEnvConfig, "service_cls": SVGService "service_cls": SVGService, "service_config_cls": SVGServiceConfig }, "primitive_skill": { "env_cls": PrimitiveSkillEnv, Loading
vagen/env/svg/__init__.py +1 −0 Original line number Diff line number Diff line from .env_config import SvgEnvConfig from .env import SVGEnv from .service import SVGService from .service_config import SVGServiceConfig No newline at end of file
vagen/env/svg/dreamsim.py +34 −0 Original line number Diff line number Diff line import torch from PIL import Image import os from dreamsim import dreamsim # Create global cache and lock, similar to DINO implementation _model_cache = {} _model_cache_lock = threading.Lock() _model_counter = 0 def get_dreamsim_model(device=None): """ Get a singleton instance of DreamSim model, using cache to avoid duplicate loading Args: device: Device to run model on Returns: DreamSimScoreCalculator: Instance of DreamSim calculator """ global _model_counter # Choose device based on availability if not specified if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" # Use device as cache key cache_key = f"dreamsim_{device}" with _model_cache_lock: if cache_key not in _model_cache: _model_counter += 1 pid = os.getpid() logging.info(f"Process {pid}: Created DreamSim model #{_model_counter} on {device}") _model_cache[cache_key] = DreamSimScoreCalculator(device=device) return _model_cache[cache_key] class DreamSimScoreCalculator: """ A wrapper class for DreamSim model to calculate similarity scores between images. Loading @@ -15,6 +48,7 @@ class DreamSimScoreCalculator: cache_dir: Cache directory for model weights device: Device to run the model on (defaults to CUDA if available, else CPU) """ cache_dir = os.path.expanduser(cache_dir) if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: Loading
vagen/env/svg/score.py +78 −48 Original line number Diff line number Diff line import numpy as np import cv2 import os from vagen.env.svg.dino import DINOScoreCalculator Loading Loading @@ -98,82 +97,113 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m return scores def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_model=None, dreamsim_model=None): def calculate_total_score_batch(gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model=None, dreamsim_model=None): """ Calculate all metrics and return a comprehensive score Calculate scores for multiple image pairs in batch mode Args: gt_im: Ground truth image gen_im: Generated image gt_code: Ground truth SVG code gen_code: Generated SVG code score_config: Dictionary containing scoring parameters - model_size: small, base, large - dino_only: Whether to use only DINO for scoring - dino_weight: Weight for DINO score - structural_weight: Weight for structural score - dreamsim_weight: Weight for DreamSim score (new) gt_images: List of ground truth images gen_images: List of generated images gt_codes: List of ground truth SVG codes gen_codes: List of generated SVG codes score_configs: List of scoring parameters dictionaries dino_model: Pre-loaded DINO model (optional) dreamsim_model: Pre-loaded DreamSim model (optional) Returns: dict: Dictionary of all scores including the total weighted score List of dictionaries containing all scores """ # Get configuration parameters with defaults model_size = score_config.get("model_size", "small") dino_only = score_config.get("dino_only", False) # Define default weights based on model size default_weights = { "small": {"dino": 3.0, "structural": 7.0, "dreamsim": 5.0}, "base": {"dino": 5.0, "structural": 5.0, "dreamsim": 5.0}, "large": {"dino": 6.0, "structural": 4.0, "dreamsim": 5.0} } batch_size = len(gt_images) if batch_size == 0: return [] # Get weights with defaults weights = { "dino": score_config.get("dino_weight", default_weights[model_size]["dino"]), "structural": score_config.get("structural_weight", default_weights[model_size]["structural"]), "dreamsim": score_config.get("dreamsim_weight", default_weights[model_size]["dreamsim"]) } # Verify all inputs have same batch size if not (len(gen_images) == len(gt_codes) == len(gen_codes) == len(score_configs) == batch_size): raise ValueError("All input lists must have the same length") # Initialize scores scores = { # Initialize results batch_results = [{ "dino_score": 0.0, "structural_score": 0.0, "dreamsim_score": 0.0, "total_score": 0.0 } } for _ in range(batch_size)] # Calculate DINO score if needed if weights["dino"] > 0: # Check if we need to calculate DINO scores need_dino = any(score_config.get("dino_weight", 5.0) > 0 for score_config in score_configs) # Check if we need to calculate DreamSim scores need_dreamsim = any(score_config.get("dreamsim_weight", 5.0) > 0 for score_config in score_configs) # Calculate DINO scores in batch if needed if need_dino: if dino_model is None: from vagen.env.svg.dino import get_dino_model # Default to small model size if not specified model_size = score_configs[0].get("model_size", "small") if score_configs else "small" dino_model = get_dino_model(model_size) scores["dino_score"] = float(dino_model.calculate_DINOv2_similarity_score(gt_im=gt_im, gen_im=gen_im)) # Calculate DreamSim score if needed if weights["dreamsim"] > 0: # Calculate all DINO scores at once using batch processing dino_scores = dino_model.calculate_batch_scores(gt_images, gen_images) # Assign scores to results for i, score in enumerate(dino_scores): batch_results[i]["dino_score"] = float(score) # Calculate DreamSim scores in batch if needed if need_dreamsim: if dreamsim_model is None: from vagen.env.svg.dreamsim import get_dreamsim_model dreamsim_model = get_dreamsim_model() scores["dreamsim_score"] = float(dreamsim_model.calculate_similarity_score(gt_im=gt_im, gen_im=gen_im)) # If DINO only mode, return only DINO score # Calculate all DreamSim scores at once using batch processing dreamsim_scores = dreamsim_model.calculate_batch_scores(gt_images, gen_images) # Assign scores to results for i, score in enumerate(dreamsim_scores): batch_results[i]["dreamsim_score"] = float(score) # Calculate structural scores and total scores for i in range(batch_size): score_config = score_configs[i] result = batch_results[i] # Check if DINO-only mode dino_only = score_config.get("dino_only", False) if dino_only: scores["total_score"] = scores["dino_score"] return scores result["total_score"] = result["dino_score"] continue # Get model size for default weights model_size = score_config.get("model_size", "small") # Define default weights based on model size default_weights = { "small": {"dino": 3.0, "structural": 7.0, "dreamsim": 5.0}, "base": {"dino": 5.0, "structural": 5.0, "dreamsim": 5.0}, "large": {"dino": 6.0, "structural": 4.0, "dreamsim": 5.0} } # Get weights with defaults weights = { "dino": score_config.get("dino_weight", default_weights[model_size]["dino"]), "structural": score_config.get("structural_weight", default_weights[model_size]["structural"]), "dreamsim": score_config.get("dreamsim_weight", default_weights[model_size]["dreamsim"]) } # Calculate structural score if needed if weights["structural"] > 0: scores["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_im, gen_im))) from vagen.env.svg.score import calculate_structural_accuracy result["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_images[i], gen_images[i]))) # Calculate weighted total score weighted_sum = ( scores["dino_score"] * weights["dino"] + scores["structural_score"] * weights["structural"] + scores["dreamsim_score"] * weights["dreamsim"] result["dino_score"] * weights["dino"] + result["structural_score"] * weights["structural"] + result["dreamsim_score"] * weights["dreamsim"] ) scores['total_score'] = max(0.0, weighted_sum) result["total_score"] = max(0.0, weighted_sum) return scores No newline at end of file return batch_results No newline at end of file
vagen/env/svg/service.py +15 −1 Original line number Diff line number Diff line Loading @@ -37,6 +37,9 @@ class SVGService(BaseService): self.model_size = self.config.model_size self.dino_model = None # Will be loaded on first use # Add DreamSim model support self.dreamsim_model = None # Will be loaded on first use if enabled # Store device for model inference self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"SVGService initialized with {self.max_workers} workers, model_size={self.model_size}, device={self.device}") Loading @@ -50,6 +53,16 @@ class SVGService(BaseService): self.dino_model = get_dino_model(self.model_size, self.device) return self.dino_model def _get_dreamsim_model(self): """ Get or initialize the DreamSim model. Uses lazy loading to avoid loading the model until needed. """ if self.dreamsim_model is None: from vagen.env.svg.dreamsim import get_dreamsim_model self.dreamsim_model = get_dreamsim_model(self.device) return self.dreamsim_model def create_environments_batch(self, ids2configs: Dict[Any, Any]) -> None: """ Create multiple SVG environments in parallel. Loading Loading @@ -181,10 +194,11 @@ class SVGService(BaseService): if valid_env_ids: # Get DINO model dino_model = self._get_dino_model() dreamsim_model = self._get_dreamsim_model() # Calculate all scores at once batch_results = calculate_total_score_batch( gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model, dreamsim_model ) # Process results directly using the index mapping Loading
