From 2cfc9be7d87ab2bdad3882762829c5f8460cdd13 Mon Sep 17 00:00:00 2001 From: Junst Date: Wed, 24 Sep 2025 14:03:15 +0900 Subject: [PATCH] Fix image shape extraction for proper dimensions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🔧 **Critical Fix**: 1. **Fixed Image Shape Extraction**: - Proper handling of [B, H, W, C] vs [B, H, W] image formats - Correct extraction of height and width dimensions - Should resolve '1024x3' dimension issue 2. **Improved Image Format Handling**: - Better detection of image tensor dimensions - Consistent height and width extraction - Should finally get correct target dimensions This should resolve the image dimension issue and allow proper concept map generation with correct target size. --- concept_attention.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/concept_attention.py b/concept_attention.py index df6d661..1da3d30 100644 --- a/concept_attention.py +++ b/concept_attention.py @@ -736,7 +736,11 @@ class ConceptAttentionProcessor: attention_outputs = self.concept_attention.extract_attention_outputs() # Create concept maps - image_shape = image.shape[-2:] # Get height and width + # Handle different image formats: [B, H, W, C] or [B, H, W] + if len(image.shape) == 4: # [B, H, W, C] + image_shape = image.shape[1:3] # Get height and width + else: # [B, H, W] + image_shape = image.shape[1:3] # Get height and width concept_maps = self.concept_attention._create_concept_maps_from_attention( concepts, concept_embeddings, image_shape )