Loading vagen/env/__init__.py +3 −1 Original line number Diff line number Diff line Loading @@ -2,6 +2,7 @@ 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.service_config import SVGServiceConfig REGISTERED_ENV = { "sokoban": { Loading @@ -22,6 +23,7 @@ REGISTERED_ENV = { "svg": { "env_cls": SVGEnv, "config_cls": SvgEnvConfig, "service_cls": SVGService "service_cls": SVGService, "service_config_cls": 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 −1 Original line number Diff line number Diff line Loading @@ -2,6 +2,7 @@ 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.service_config import SVGServiceConfig REGISTERED_ENV = { "sokoban": { Loading @@ -22,6 +23,7 @@ REGISTERED_ENV = { "svg": { "env_cls": SVGEnv, "config_cls": SvgEnvConfig, "service_cls": SVGService "service_cls": SVGService, "service_config_cls": 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
