Update controlnet_module.py

controlnet_module.py

@@ -159,34 +159,33 @@ class ControlNetProcessor:
             print(f"🔍 Original_sizes parameter (in inputs): {inputs.get('original_sizes', 'NOT FOUND')}")
             print(f"🔍 Reshaped_input_sizes parameter (in inputs): {inputs.get('reshaped_input_sizes', 'NOT FOUND')}")
 
-
+            # 1. Maske auswählen
+            single_mask = outputs.pred_masks[:, :, 0, :, :]  # Shape: [1, 1, 256, 256]
 
             # Debug: Vollständige Dimensionsanalyse
             print(f"🔍 outputs.pred_masks shape: {outputs.pred_masks.shape}")
             print(f"🔍 outputs.pred_masks dimensions: {outputs.pred_masks.dim()}")
 
-            # Nach der Auswahl
-            single_mask = outputs.pred_masks[:, :, 0, :, :]
-            print(f"🔍 single_mask shape: {single_mask.shape}")
-            print(f"🔍 single_mask dimensions: {single_mask.dim()}")
-
-            # Wichtig: Ist es wirklich [1, 1, 256, 256] oder [1, 256, 256]?
+            # 2. KRITISCHE KORREKTUR: Stelle sicher, dass es 4D bleibt
+            # Wenn single_mask 3D ist ([1, 256, 256]), füge Channel-Dimension hinzu
             if single_mask.dim() == 3:
-                print("⚠️  WARNUNG: Maske hat nur 3 Dimensionen! Korrigiere...")
-                single_mask = single_mask.unsqueeze(1)  # Fügt Channel-Dimension hinzu
+                print("⚠️  Maske ist 3D, füge Channel-Dimension hinzu")
+                single_mask = single_mask.unsqueeze(1)  # Wird zu [1, 1, 256, 256]
 
+            print(f"🔍 Final mask shape before post_process: {single_mask.shape}")
+            print(f"🔍 Final mask dimensions: {single_mask.dim()}")
 
-            # 2. Konvertiere die Größen-Parameter zu CPU und als Python-Liste
+            # 3. Größen-Parameter konvertieren
             original_size = inputs['original_sizes'].cpu().numpy().tolist()[0]  # [512, 512]
             input_size = inputs['reshaped_input_sizes'].cpu().numpy().tolist()[0]  # [512, 512]
 
-            # 3. Korrekter Aufruf mit angepassten Parametern
+            # 4. Maske verarbeiten (WICHTIG: original_size muss Tuple/Liste sein)
             mask = self.sam_processor.post_process_masks(
-                single_mask,                    # Nur eine Maske
-                original_sizes=[original_size], # Als Python-Liste
-                reshaped_input_sizes=[input_size] # Als Python-Liste
-            )[0][0]  # [batch][mask]
-
+                single_mask,
+                original_sizes=[original_size],  # [512, 512] als Liste in einer Liste
+                reshaped_input_sizes=[input_size]  # [512, 512] als Liste in einer Liste
+            )[0][0]
+   
 
             # Sigmoid für Wahrscheinlichkeiten, dann Schwellenwert
             mask = mask.sigmoid().cpu().numpy()