Loading vagen/env/README.md +6 −0 Original line number Diff line number Diff line Loading @@ -21,3 +21,9 @@ python vagen/env/navigation/startx.py 1 pip install gymnasium pip install "gymnasium[toy-text]" ``` ### SVG ``` # Download dataset from huggingface ``` No newline at end of file vagen/env/__init__.py +5 −0 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanConfig from .frozenlake import FrozenLakeEnv,FrozenLakeConfig from .navigation import NavigationEnv, NavigationConfig from .svg import SVGEnv, SVGConfig REGISTERED_ENV = { "sokoban": { "env_cls": SokobanEnv, Loading @@ -13,5 +14,9 @@ REGISTERED_ENV = { "navigation": { "env_cls": NavigationEnv, "config_cls": NavigationConfig }, "svg": { "env_cls": SVGEnv, "config_cls": SVGConfig } } No newline at end of file vagen/env/svg/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .config import SVGConfig from .env import SVGEnv No newline at end of file vagen/env/svg/config.py 0 → 100644 +79 −0 Original line number Diff line number Diff line from vagen.env.base_config import BaseConfig from dataclasses import dataclass, fields, field from typing import Optional, List, Union, Dict @dataclass class SVGConfig(BaseConfig): """Configuration for the SVG environment""" dataset_name: str = "starvector/svg-emoji-simple" data_dir: str = "vagen/env/svg/data" seed: int = 16 split: str = "train" # Score configuration model_size: str = "small" # 'small', 'base', or 'large' dino_only: bool = False dino_weight: Optional[float] = None structural_weight: Optional[float] = None color_weight: Optional[float] = None code_weight: Optional[float] = None # Reward configuration format_reward: float = 0.5 format_penalty: float = 0.0 # Analysis mode for logging analysis_mode: bool = False def config_id(self) -> str: """Generate a unique identifier for this configuration""" id_fields = [ "dataset_name", "model_size", "dino_only", "format_reward", "format_penalty" ] id_str = ",".join([f"{field.name}={getattr(self, field.name)}" for field in fields(self) if field.name in id_fields]) # Add optional fields if they're set optional_fields = ["dino_weight", "structural_weight", "color_weight", "code_weight"] for field_name in optional_fields: value = getattr(self, field_name) if value is not None: id_str += f",{field_name}={value}" return f"SVGConfig({id_str})" def get_score_config(self) -> Dict: """Get the score configuration dictionary""" score_config = { "model_size": self.model_size, "dino_only": self.dino_only, } # Add optional weights if set if self.dino_weight is not None: score_config["dino_weight"] = self.dino_weight if self.structural_weight is not None: score_config["structural_weight"] = self.structural_weight if self.color_weight is not None: score_config["color_weight"] = self.color_weight if self.code_weight is not None: score_config["code_weight"] = self.code_weight return score_config if __name__ == "__main__": # Example usage config = SVGConfig( dataset_name="starvector/svg-emoji-simple", data_dir="data/svg", model_size="small", dino_only=False, dino_weight=5.0 ) print(config.config_id()) print(config.get_score_config()) No newline at end of file vagen/env/svg/dino.py 0 → 100644 +124 −0 Original line number Diff line number Diff line import torch from torch.utils.data import DataLoader from tqdm import tqdm from transformers import AutoModel, AutoImageProcessor from PIL import Image import torch.nn as nn # @TODO clean codes of this section class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self): self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count class BaseMetric: def __init__(self): self.meter = AverageMeter() def reset(self): self.meter.reset() def calculate_score(self, batch, update=True): """ Batch: {"gt_im": [PIL Image], "gen_im": [Image]} """ values = [] batch_size = len(next(iter(batch.values()))) for index in tqdm(range(batch_size)): kwargs = {} for key in ["gt_im", "gen_im", "gt_svg", "gen_svg", "caption"]: if key in batch: kwargs[key] = batch[key][index] try: measure = self.metric(**kwargs) except Exception as e: print("Error calculating metric: {}".format(e)) continue if math.isnan(measure): continue values.append(measure) if not values: print("No valid values found for metric calculation.") return float("nan") score = sum(values) / len(values) if update: self.meter.update(score, len(values)) return self.meter.avg, values else: return score, values def metric(self, **kwargs): """ This method should be overridden by subclasses to provide the specific metric computation. """ raise NotImplementedError("The metric method must be implemented by subclasses.") def get_average_score(self): return self.meter.avg class DINOScoreCalculator(BaseMetric): #@TODO how to make sure DINO always on GPU? check how ray is deliver gpu resources def __init__(self, config=None, model_size='large', device='cuda'): super().__init__() self.class_name = self.__class__.__name__ self.config = config self.model_size = model_size self.model, self.processor = self.get_DINOv2_model(model_size) device = device if torch.cuda.is_available() else "cpu" self.model = self.model.to(device) self.device = device self.metric = self.calculate_DINOv2_similarity_score def get_DINOv2_model(self, model_size): if model_size == "small": model_size = "facebook/dinov2-small" elif model_size == "base": model_size = "facebook/dinov2-base" elif model_size == "large": model_size = "facebook/dinov2-large" else: raise ValueError(f"model_size should be either 'small', 'base' or 'large', got {model_size}") return AutoModel.from_pretrained(model_size), AutoImageProcessor.from_pretrained(model_size) def process_input(self, image, processor): if isinstance(image, str): image = Image.open(image) if isinstance(image, Image.Image): with torch.no_grad(): inputs = processor(images=image, return_tensors="pt").to(self.device) outputs = self.model(**inputs) features = outputs.last_hidden_state.mean(dim=1) elif isinstance(image, torch.Tensor): features = image.unsqueeze(0) if image.dim() == 1 else image else: raise ValueError("Input must be a file path, PIL Image, or tensor of features") return features def calculate_DINOv2_similarity_score(self, **kwargs): image1 = kwargs.get('gt_im') image2 = kwargs.get('gen_im') features1 = self.process_input(image1, self.processor) features2 = self.process_input(image2, self.processor) cos = nn.CosineSimilarity(dim=1) sim = cos(features1, features2).item() sim = (sim + 1) / 2 return sim No newline at end of file Loading
vagen/env/README.md +6 −0 Original line number Diff line number Diff line Loading @@ -21,3 +21,9 @@ python vagen/env/navigation/startx.py 1 pip install gymnasium pip install "gymnasium[toy-text]" ``` ### SVG ``` # Download dataset from huggingface ``` No newline at end of file
vagen/env/__init__.py +5 −0 Original line number Diff line number Diff line from .sokoban import SokobanEnv,SokobanConfig from .frozenlake import FrozenLakeEnv,FrozenLakeConfig from .navigation import NavigationEnv, NavigationConfig from .svg import SVGEnv, SVGConfig REGISTERED_ENV = { "sokoban": { "env_cls": SokobanEnv, Loading @@ -13,5 +14,9 @@ REGISTERED_ENV = { "navigation": { "env_cls": NavigationEnv, "config_cls": NavigationConfig }, "svg": { "env_cls": SVGEnv, "config_cls": SVGConfig } } No newline at end of file
vagen/env/svg/__init__.py 0 → 100644 +2 −0 Original line number Diff line number Diff line from .config import SVGConfig from .env import SVGEnv No newline at end of file
vagen/env/svg/config.py 0 → 100644 +79 −0 Original line number Diff line number Diff line from vagen.env.base_config import BaseConfig from dataclasses import dataclass, fields, field from typing import Optional, List, Union, Dict @dataclass class SVGConfig(BaseConfig): """Configuration for the SVG environment""" dataset_name: str = "starvector/svg-emoji-simple" data_dir: str = "vagen/env/svg/data" seed: int = 16 split: str = "train" # Score configuration model_size: str = "small" # 'small', 'base', or 'large' dino_only: bool = False dino_weight: Optional[float] = None structural_weight: Optional[float] = None color_weight: Optional[float] = None code_weight: Optional[float] = None # Reward configuration format_reward: float = 0.5 format_penalty: float = 0.0 # Analysis mode for logging analysis_mode: bool = False def config_id(self) -> str: """Generate a unique identifier for this configuration""" id_fields = [ "dataset_name", "model_size", "dino_only", "format_reward", "format_penalty" ] id_str = ",".join([f"{field.name}={getattr(self, field.name)}" for field in fields(self) if field.name in id_fields]) # Add optional fields if they're set optional_fields = ["dino_weight", "structural_weight", "color_weight", "code_weight"] for field_name in optional_fields: value = getattr(self, field_name) if value is not None: id_str += f",{field_name}={value}" return f"SVGConfig({id_str})" def get_score_config(self) -> Dict: """Get the score configuration dictionary""" score_config = { "model_size": self.model_size, "dino_only": self.dino_only, } # Add optional weights if set if self.dino_weight is not None: score_config["dino_weight"] = self.dino_weight if self.structural_weight is not None: score_config["structural_weight"] = self.structural_weight if self.color_weight is not None: score_config["color_weight"] = self.color_weight if self.code_weight is not None: score_config["code_weight"] = self.code_weight return score_config if __name__ == "__main__": # Example usage config = SVGConfig( dataset_name="starvector/svg-emoji-simple", data_dir="data/svg", model_size="small", dino_only=False, dino_weight=5.0 ) print(config.config_id()) print(config.get_score_config()) No newline at end of file
vagen/env/svg/dino.py 0 → 100644 +124 −0 Original line number Diff line number Diff line import torch from torch.utils.data import DataLoader from tqdm import tqdm from transformers import AutoModel, AutoImageProcessor from PIL import Image import torch.nn as nn # @TODO clean codes of this section class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self): self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count class BaseMetric: def __init__(self): self.meter = AverageMeter() def reset(self): self.meter.reset() def calculate_score(self, batch, update=True): """ Batch: {"gt_im": [PIL Image], "gen_im": [Image]} """ values = [] batch_size = len(next(iter(batch.values()))) for index in tqdm(range(batch_size)): kwargs = {} for key in ["gt_im", "gen_im", "gt_svg", "gen_svg", "caption"]: if key in batch: kwargs[key] = batch[key][index] try: measure = self.metric(**kwargs) except Exception as e: print("Error calculating metric: {}".format(e)) continue if math.isnan(measure): continue values.append(measure) if not values: print("No valid values found for metric calculation.") return float("nan") score = sum(values) / len(values) if update: self.meter.update(score, len(values)) return self.meter.avg, values else: return score, values def metric(self, **kwargs): """ This method should be overridden by subclasses to provide the specific metric computation. """ raise NotImplementedError("The metric method must be implemented by subclasses.") def get_average_score(self): return self.meter.avg class DINOScoreCalculator(BaseMetric): #@TODO how to make sure DINO always on GPU? check how ray is deliver gpu resources def __init__(self, config=None, model_size='large', device='cuda'): super().__init__() self.class_name = self.__class__.__name__ self.config = config self.model_size = model_size self.model, self.processor = self.get_DINOv2_model(model_size) device = device if torch.cuda.is_available() else "cpu" self.model = self.model.to(device) self.device = device self.metric = self.calculate_DINOv2_similarity_score def get_DINOv2_model(self, model_size): if model_size == "small": model_size = "facebook/dinov2-small" elif model_size == "base": model_size = "facebook/dinov2-base" elif model_size == "large": model_size = "facebook/dinov2-large" else: raise ValueError(f"model_size should be either 'small', 'base' or 'large', got {model_size}") return AutoModel.from_pretrained(model_size), AutoImageProcessor.from_pretrained(model_size) def process_input(self, image, processor): if isinstance(image, str): image = Image.open(image) if isinstance(image, Image.Image): with torch.no_grad(): inputs = processor(images=image, return_tensors="pt").to(self.device) outputs = self.model(**inputs) features = outputs.last_hidden_state.mean(dim=1) elif isinstance(image, torch.Tensor): features = image.unsqueeze(0) if image.dim() == 1 else image else: raise ValueError("Input must be a file path, PIL Image, or tensor of features") return features def calculate_DINOv2_similarity_score(self, **kwargs): image1 = kwargs.get('gt_im') image2 = kwargs.get('gen_im') features1 = self.process_input(image1, self.processor) features2 = self.process_input(image2, self.processor) cos = nn.CosineSimilarity(dim=1) sim = cos(features1, features2).item() sim = (sim + 1) / 2 return sim No newline at end of file
