Update controlnet_module.py

controlnet_module.py

@@ -134,14 +134,12 @@ class ControlNetProcessor:
 
             # 3. Vorbereiten der Eingabe für SAM2
             # BBox im Format [x_min, y_min, x_max, y_max] erstellen
-            # ACHTUNG: SAM2 erwartet Boxen in diesem Format
-            # Zeilen in der Funktion anpassen:
-            
-            input_boxes = [[[x1, y1, x2, y2]]] #Dreifach verschachtelt
+            # Dreifach verschachteltes Format: [[[x1, y1, x2, y2]]]
+            input_boxes = [[[x1, y1, x2, y2]]]
 
-            # Bild mit dem Processor vorverarbeiten
+            # Original-Bild und BBox-Koordinaten zur Segmentierung vorverarbeiten
             inputs = self.sam_processor(
-                image_np,
+                image_np,  
                 input_boxes=input_boxes,
                 return_tensors="pt"
             ).to(self.device)
@@ -151,56 +149,43 @@ class ControlNetProcessor:
             with torch.no_grad():
                 outputs = self.sam_model(**inputs)
 
-
             # DEBUG: Dimensionen prüfen
-            print(f"🔍 Original image size: {image.size}")  # z.B. (1024, 768)
-            print(f"🔍 Processed image size: {inputs['pixel_values'].shape}")  # z.B. torch.Size([1, 3, 1024, 1024])
-            print(f"🔍 Output masks shape: {outputs.pred_masks.shape}")  # KRITISCH: z.B. torch.Size([1, 2, 256, 256])
-            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')}")
+            print(f"🔍 Original image size: {image.size}")
+            print(f"🔍 Processed image size: {inputs['pixel_values'].shape}")
+            print(f"🔍 Output masks shape: {outputs.pred_masks.shape}")
 
-            # 1. Maske auswählen
+            # 5. Maske auswählen (erste Maske der ersten Batch-Dimension)
             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()}")
-
-            # 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("⚠️  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()}")
-
-            # 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]
-
-            # 4. Maske verarbeiten (WICHTIG: original_size muss Tuple/Liste sein)
-            mask = self.sam_processor.post_process_masks(
-                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]
-   
+            print(f"🔍 Single mask shape: {single_mask.shape}")
+            print(f"🔍 Single mask dimensions: {single_mask.dim()}")
 
-            # Sigmoid für Wahrscheinlichkeiten, dann Schwellenwert
-            mask = mask.sigmoid().cpu().numpy()
-            mask_array = (mask > 0.5).astype(np.uint8) * 255
+            # 6. KRITISCHE KORREKTUR: Direkte Skalierung statt post_process_masks
+            import torch.nn.functional as F
+            
+            # Skaliere die 256x256 Rohmaske direkt auf Ihre Zielgröße (image.height, image.width)
+            final_mask = F.interpolate(
+                single_mask,  # Direkt die 256x256 Rohmaske verwenden
+                size=(image.height, image.width),  # Direkt auf Ihre Zielgröße skalieren
+                mode='bilinear',
+                align_corners=False
+            ).squeeze()  # Entferne Batch- und Channel-Dimensionen
 
-            # 6. Zu PIL Image konvertieren und auf Originalgröße bringen
-            mask = Image.fromarray(mask_array.squeeze()).convert("L")
-            mask = mask.resize((width, height), Image.Resampling.NEAREST)
+            print(f"🔍 Final mask shape after interpolation: {final_mask.shape}")
+
+            # 7. In NumPy konvertieren und Schwellenwert anwenden
+            mask_np = final_mask.sigmoid().cpu().numpy()
+            mask_array = (mask_np > 0.5).astype(np.uint8) * 255
 
-            # 7. Kanten glätten für natürlichere Übergänge
+            # 8. Zu PIL Image konvertieren
+            mask = Image.fromarray(mask_array.squeeze()).convert("L")
+            
+            # 9. Kanten glätten für natürlichere Übergänge
             mask_array = np.array(mask)
             mask_array = self._smooth_mask(mask_array, blur_radius=2)
             mask = Image.fromarray(mask_array).convert("L")
 
-            # 8. Modus-spezifische Anpassung (Invertierung)
+            # 10. Modus-spezifische Anpassung (Invertierung)
             if mode == "environment_change":
                 # MODUS 1: Umgebung ändern - Objekt schwarz (erhalten)
                 mask = Image.eval(mask, lambda x: 255 - x)
@@ -211,7 +196,7 @@ class ControlNetProcessor:
 
             print(f"✅ SAM 2: Präzise Maske erstellt ({mask.size})")
             return mask
-
+    
         except Exception as e:
             print(f"⚠️ SAM 2 Fehler (Transformers API): {str(e)[:200]}")
             print(f"🔍 SAM 2 Model dtype: {self.sam_model.dtype}")
@@ -221,7 +206,8 @@ class ControlNetProcessor:
             traceback.print_exc()
             print("ℹ️ Fallback auf rechteckige Maske")
             return self._create_rectangular_mask(image, bbox_coords, mode)
-
+            
+   
     def _create_rectangular_mask(self, image, bbox_coords, mode):
         """Fallback: Erstellt rechteckige Maske"""
         from PIL import ImageDraw