add dreamsim weight (2a7a670d) · Commits · 张泽凯 / VAGEN

vagen/env/svg/dreamsim.py

0 → 100644

+100 −0

Original line number	Diff line number	Diff line
		import torch
		from PIL import Image
		from dreamsim import dreamsim

		class DreamSimScoreCalculator:
		"""
		A wrapper class for DreamSim model to calculate similarity scores between images.
		"""
		def __init__(self, pretrained=True, cache_dir="~/.cache", device=None):
		"""
		Initialize DreamSim model.

		Args:
		pretrained: Whether to use pretrained model
		cache_dir: Cache directory for model weights
		device: Device to run the model on (defaults to CUDA if available, else CPU)
		"""
		if device is None:
		self.device = "cuda" if torch.cuda.is_available() else "cpu"
		else:
		self.device = device

		# Load model and preprocessor
		self.model, self.preprocess = dreamsim(pretrained=pretrained, cache_dir=cache_dir)
		self.model = self.model.to(self.device)

		def calculate_similarity_score(self, gt_im, gen_im):
		"""
		Calculate similarity score between ground truth and generated images.

		Args:
		gt_im: Ground truth PIL Image
		gen_im: Generated PIL Image

		Returns:
		float: Similarity score (1 - distance, normalized to [0, 1])
		"""
		# Preprocess images
		img1 = self.preprocess(gt_im)
		img2 = self.preprocess(gen_im)

		# Move to device if necessary
		img1 = img1.to(self.device)
		img2 = img2.to(self.device)

		# Calculate distance (lower is better)
		with torch.no_grad():
		distance = self.model(img1, img2).item()

		# Convert distance to similarity score (1 - normalized distance)
		# DreamSim usually outputs values in range [0, 1] where lower means more similar
		# We invert it so that higher means more similar (1 = identical)
		similarity = 1.0 - min(1.0, max(0.0, distance))

		return similarity

		def calculate_batch_scores(self, gt_images, gen_images):
		"""
		Calculate similarity scores for a batch of image pairs.

		Args:
		gt_images: List of ground truth PIL Images
		gen_images: List of generated PIL Images

		Returns:
		List[float]: List of similarity scores
		"""
		# Preprocess all images
		gt_processed = [self.preprocess(img) for img in gt_images]
		gen_processed = [self.preprocess(img) for img in gen_images]

		scores = []
		# Process each pair
		for gt, gen in zip(gt_processed, gen_processed):
		# Move to device
		gt = gt.to(self.device)
		gen = gen.to(self.device)

		# Calculate distance
		with torch.no_grad():
		distance = self.model(gt, gen).item()

		# Convert to similarity score
		similarity = 1.0 - min(1.0, max(0.0, distance))
		scores.append(similarity)

		return scores

		# Helper function to get or initialize DreamSim model
		def get_dreamsim_model(device=None):
		"""
		Get an instance of DreamSim model.

		Args:
		device: Device to run model on

		Returns:
		DreamSimScoreCalculator: Instance of DreamSim calculator
		"""
		return DreamSimScoreCalculator(device=device)
		No newline at end of file

vagen/env/svg/env.py

+0 −2

Original line number	Diff line number	Diff line
		@@ -224,8 +224,6 @@ class SVGEnv(BaseEnv):
		"multi_modal_data": multi_modal_data,
		}

		def set_dino_model(self, model):
		self.dino_model = model

		if __name__ == "__main__":
		config = SvgEnvConfig()

vagen/env/svg/env_config.py

+4 −1

Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@ class SvgEnvConfig(BaseEnvConfig):
		dino_only: bool = False
		dino_weight: Optional[float] = None
		structural_weight: Optional[float] = None
		dreamsim_weight: Optional[float] = None
		# Reward configuration
		format_reward: float = 0.5
		format_penalty: float = 0.0
		@@ -34,7 +35,7 @@ class SvgEnvConfig(BaseEnvConfig):
		if field.name in id_fields])

		# Add optional fields if they're set
		optional_fields = ["dino_weight", "structural_weight"]
		optional_fields = ["dino_weight", "structural_weight", "dreamsim_weight"]
		for field_name in optional_fields:
		value = getattr(self, field_name)
		if value is not None:
		@@ -54,6 +55,8 @@ class SvgEnvConfig(BaseEnvConfig):
		score_config["dino_weight"] = self.dino_weight
		if self.structural_weight is not None:
		score_config["structural_weight"] = self.structural_weight
		if self.dreamsim_weight is not None:
		score_config["dreamsim_weight"] = self.dreamsim_weight

		return score_config

vagen/env/svg/prompt.py

+2 −14

Original line number	Diff line number	Diff line
		@@ -5,22 +5,10 @@ system_prompt = """You are a precise SVG code generator. You will be given an im
		SVG Quick Guide
		Goal: Transform the provided image into precise SVG code that replicates the image.

		Basic SVG Elements:
		- <rect> for rectangles and squares
		- <circle> for circles
		- <ellipse> for ellipses
		- <line> for straight lines
		- <polyline> for connected lines
		- <polygon> for closed shapes
		- <path> for complex shapes and curves
		- <text> for text elements
		- <g> for grouping elements

		Process:
		1. First analyze the image carefully, identifying distinct visual elements
		2. Break down complex shapes into basic SVG elements when possible
		3. Identify colors, dimensions, positions, and relationships between elements
		4. Generate accurate SVG code that reproduces the image
		2. Identify colors, dimensions, positions, and relationships between elements
		3. Generate accurate SVG code that reproduces the image

		Rewards:
		- Overall visual similarity: +5.0

vagen/env/svg/score.py

+70 −67

Original line number	Diff line number	Diff line
		@@ -18,7 +18,7 @@ def calculate_structural_accuracy(gt_im, gen_im):
		return intersection / union if union > 0 else 0


		def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_model=None):
		def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_model=None, dreamsim_model=None):
		"""
		Calculate all metrics and return a comprehensive score

		@@ -32,7 +32,9 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m
		- 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)
		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
		@@ -43,21 +45,23 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m

		# Define default weights based on model size
		default_weights = {
		"small": {"dino": 3.0, "structural": 7.0},
		"base": {"dino": 5.0, "structural": 5.0},
		"large": {"dino": 6.0, "structural": 4.0}
		"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"])
		"structural": score_config.get("structural_weight", default_weights[model_size]["structural"]),
		"dreamsim": score_config.get("dreamsim_weight", default_weights[model_size]["dreamsim"])
		}

		# Initialize scores
		scores = {
		"dino_score": 0.0,
		"structural_score": 0.0,
		"dreamsim_score": 0.0,
		"total_score": 0.0
		}

		@@ -68,6 +72,13 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m
		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:
		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
		if dino_only:
		scores["total_score"] = scores["dino_score"]
		@@ -80,97 +91,89 @@ def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_m
		# Calculate weighted total score
		weighted_sum = (
		scores["dino_score"] * weights["dino"] +
		scores["structural_score"] * weights["structural"]
		scores["structural_score"] * weights["structural"] +
		scores["dreamsim_score"] * weights["dreamsim"]
		)
		scores['total_score'] = max(0.0, weighted_sum)

		return scores

		def calculate_total_score_batch(gt_images, gen_images, gt_codes, gen_codes, score_configs, dino_model=None):
		def calculate_total_score(gt_im, gen_im, gt_code, gen_code, score_config, dino_model=None, dreamsim_model=None):
		"""
		Calculate scores for multiple image pairs in batch mode
		Calculate all metrics and return a comprehensive score

		Args:
		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
		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)
		dino_model: Pre-loaded DINO model (optional)
		dreamsim_model: Pre-loaded DreamSim model (optional)

		Returns:
		List of dictionaries containing all scores
		dict: Dictionary of all scores including the total weighted score
		"""
		batch_size = len(gt_images)
		if batch_size == 0:
		return []
		# Get configuration parameters with defaults
		model_size = score_config.get("model_size", "small")
		dino_only = score_config.get("dino_only", False)

		# 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")
		# 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}
		}

		# Initialize results
		batch_results = [{
		# 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"])
		}

		# Initialize scores
		scores = {
		"dino_score": 0.0,
		"structural_score": 0.0,
		"dreamsim_score": 0.0,
		"total_score": 0.0
		} for _ in range(batch_size)]

		# 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)
		}

		# Calculate DINO scores in batch if needed
		if need_dino:
		# Calculate DINO score if needed
		if weights["dino"] > 0:
		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 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 structural scores and total scores
		for i in range(batch_size):
		score_config = score_configs[i]
		result = batch_results[i]
		# Calculate DreamSim score if needed
		if weights["dreamsim"] > 0:
		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))

		# Check if DINO-only mode
		dino_only = score_config.get("dino_only", False)
		# If DINO only mode, return only DINO score
		if dino_only:
		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},
		"base": {"dino": 5.0, "structural": 5.0},
		"large": {"dino": 6.0, "structural": 4.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"])
		}
		scores["total_score"] = scores["dino_score"]
		return scores

		# Calculate structural score if needed
		if weights["structural"] > 0:
		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])))
		scores["structural_score"] = max(0.0, float(calculate_structural_accuracy(gt_im, gen_im)))

		# Calculate weighted total score
		weighted_sum = (
		result["dino_score"] * weights["dino"] +
		result["structural_score"] * weights["structural"]
		scores["dino_score"] * weights["dino"] +
		scores["structural_score"] * weights["structural"] +
		scores["dreamsim_score"] * weights["dreamsim"]
		)
		result["total_score"] = max(0.0, weighted_sum)
		scores['total_score'] = max(0.0, weighted_sum)

		return batch_results
		No newline at end of file
		return scores
		No newline at end of file

Admin message