Update controlnet_module.py

controlnet_module.py

@@ -6,7 +6,8 @@ import random
 import cv2
 import numpy as np
 import gradio as gr
-from segment_anything import sam_model_registry, SamPredictor
+# WICHTIG: Importiere die neuen SAM2-Klassen aus Transformers
+from transformers import Sam2Model, Sam2Processor
 
 
 class ControlNetProgressCallback:
@@ -33,100 +34,57 @@ class ControlNetProcessor:
         self.pose_detector = None
         self.midas_model = None
         self.midas_transform = None
-        self.sam_predictor = None
+        # Ändere die Variablennamen für die neue API
+        self.sam_processor = None
+        self.sam_model = None
         self.sam_initialized = False
 
-
-    # In controlnet_module.py - Ersetze die _lazy_load_sam() Funktion
-from transformers import Sam2Model, Sam2Processor
-
-
-    # In controlnet_module.py - Ersetze die _lazy_load_sam() Funktion
-from transformers import Sam2Model, Sam2Processor
-
-def _lazy_load_sam(self):
-    if self.sam_initialized:
-        return True
-        
-    try:
-        print("🔄 Lade SAM 2 über 🤗 Transformers...")
-        
-        model_id = "facebook/sam2-hiera-tiny"  # Dieser Pfad ist korrekt
-        self.sam_processor = Sam2Processor.from_pretrained(model_id)
-        self.sam_model = Sam2Model.from_pretrained(model_id).to(self.device)
-        
-        self.sam_initialized = True
-        print("✅ SAM 2 erfolgreich geladen (via Transformers)")
-        return True
-        
-    except Exception as e:
-        print(f"❌ Fehler beim Laden von SAM 2: {e}")
-        self.sam_initialized = True
-        return False
-
-
-def _lazy_load_sam(self):
-    if self.sam_initialized:
-        return True
-        
-    try:
-        print("🔄 Lade SAM 2 über 🤗 Transformers...")
-        
-        model_id = "facebook/sam2-hiera-tiny"  # Dieser Pfad ist korrekt
-        self.sam_processor = Sam2Processor.from_pretrained(model_id)
-        self.sam_model = Sam2Model.from_pretrained(model_id).to(self.device)
-        
-        self.sam_initialized = True
-        print("✅ SAM 2 erfolgreich geladen (via Transformers)")
-        return True
-        
-    except Exception as e:
-        print(f"❌ Fehler beim Laden von SAM 2: {e}")
-        self.sam_initialized = True
-        return False
-        
     def _lazy_load_sam(self):
-        """Lazy Loading von SAM 2 Tiny - Optimiert für Hugging Face Spaces"""
+        """Lazy Loading von SAM 2 über 🤗 Transformers API"""
         if self.sam_initialized:
             return True
-            
+
         try:
-            print("🔄 Lade SAM 2 Tiny von Hugging Face Hub...")
-            
-            # KORRIGIERT: Nur der Hugging Face Model-ID Pfad
+            print("🔄 Lade SAM 2 über 🤗 Transformers...")
+
+            # Die korrekte Modell-ID für SAM 2 Tiny
             model_id = "facebook/sam2-hiera-tiny"
-            
-            # SAM 2 Modell direkt von Hugging Face laden
-            sam = sam_model_registry["sam2_hiera_tiny"](checkpoint=model_id)
-            sam.to(self.device)
-            self.sam_predictor = SamPredictor(sam)
-            
+
+            # Lade Processor und Modell mit der neuen API
+            self.sam_processor = Sam2Processor.from_pretrained(model_id)
+            self.sam_model = Sam2Model.from_pretrained(model_id).to(self.device)
+            self.sam_model.eval()  # Setze Modell in Evaluierungsmodus
+
             self.sam_initialized = True
-            print(f"✅ SAM 2 ({model_id}) erfolgreich geladen")
+            print("✅ SAM 2 erfolgreich geladen (via Transformers)")
             return True
-            
+
         except Exception as e:
-            print(f"❌ SAM 2 konnte nicht geladen werden: {str(e)[:100]}")
-            print("ℹ️ Verwende rechteckige Masken als Fallback")
-            self.sam_predictor = None
+            print(f"❌ Fehler beim Laden von SAM 2: {str(e)[:200]}")
             self.sam_initialized = True  # Verhindert weitere Ladeversuche
             return False
-    
+
     def _validate_bbox(self, image, bbox_coords):
         """Validiert und korrigiert BBox-Koordinaten"""
         width, height = image.size
-        x1, y1, x2, y2 = bbox_coords
-        
+
+        # Extrahiere Koordinaten - unterstützt beide Formate
+        if isinstance(bbox_coords, (list, tuple)) and len(bbox_coords) == 4:
+            x1, y1, x2, y2 = bbox_coords
+        else:
+            # Für den Fall, dass Koordinaten einzeln übergeben werden
+            x1, y1, x2, y2 = bbox_coords
+
         # Stelle sicher, dass x1 <= x2 und y1 <= y2
         x1, x2 = min(x1, x2), max(x1, x2)
         y1, y2 = min(y1, y2), max(y1, y2)
-        
+
         # Begrenze auf Bildgrenzen
         x1 = max(0, min(x1, width - 1))
         y1 = max(0, min(y1, height - 1))
         x2 = max(0, min(x2, width - 1))
         y2 = max(0, min(y2, height - 1))
-        
+
         # Stelle sicher, dass BBox gültig ist
         if x2 - x1 < 10 or y2 - y1 < 10:
             # Fallback auf sinnvolle Größe
@@ -135,110 +93,121 @@ def _lazy_load_sam(self):
             y1 = max(0, height/2 - size/2)
             x2 = min(width, width/2 + size/2)
             y2 = min(height, height/2 + size/2)
-        
+
         return int(x1), int(y1), int(x2), int(y2)
-    
+
     def _smooth_mask(self, mask_array, blur_radius=3):
-        """Glättet die Maske für bessere Übergänge (5-Pixel Randbereich)"""
+        """Glättet die Maske für bessere Übergänge"""
         try:
-            # Gaussian Blur für weiche Kanten - nur der Randbereich wird beeinflusst
             if blur_radius > 0:
-                mask_array = cv2.GaussianBlur(mask_array, 
-                                            (blur_radius*2+1, blur_radius*2+1), 
-                                            0)
-            
+                # Verwende median blur für bessere Kantenerhaltung als Gaussian
+                mask_array = cv2.medianBlur(mask_array, blur_radius*2+1)
             return mask_array
-        except:
+        except Exception as e:
+            print(f"⚠️ Fehler beim Glätten der Maske: {e}")
             return mask_array
-    
+
     def create_sam_mask(self, image, bbox_coords, mode):
         """
-        Erstellt präzise Maske mit SAM 2 (transparent für Benutzer)
+        Erstellt präzise Maske mit SAM 2 (via 🤗 Transformers API)
         Gibt PIL Image in L-Modus zurück (0=schwarz=erhalten, 255=weiß=verändern)
         """
         try:
-            # Lade SAM bei Bedarf (automatisch für Hugging Face Spaces)
+            # 1. SAM2 laden (falls noch nicht geschehen)
             if not self.sam_initialized:
                 self._lazy_load_sam()
-            
-            # Fallback wenn SAM nicht verfügbar
-            if self.sam_predictor is None:
+
+            if self.sam_model is None or self.sam_processor is None:
+                print("⚠️ SAM 2 Model nicht verfügbar, verwende Fallback")
                 return self._create_rectangular_mask(image, bbox_coords, mode)
-            
-            # Validiere BBox
+
+            # 2. Validiere BBox und konvertiere Bild
             x1, y1, x2, y2 = self._validate_bbox(image, bbox_coords)
-            
-            # Konvertiere zu numpy array (RGB)
+            width, height = image.size
+
+            # Konvertiere zu numpy array (RGB) - für SAM2 Processor
             image_np = np.array(image.convert("RGB"))
-            
-            # SAM vorbereiten
-            try:
-                self.sam_predictor.set_image(image_np)
-            except Exception as e:
-                print(f"⚠️ SAM set_image Fehler: {e}")
-                return self._create_rectangular_mask(image, bbox_coords, mode)
-            
-            # BBox für SAM formatieren
-            input_box = np.array([x1, y1, x2, y2])
-            
+
+            # 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
+            input_boxes = [[x1, y1, x2, y2]]
+
+            # Bild mit dem Processor vorverarbeiten
+            inputs = self.sam_processor(
+                image_np,
+                input_boxes=[input_boxes],  # WICHTIG: Liste von Box-Listen
+                return_tensors="pt"
+            ).to(self.device)
+
+            # 4. Vorhersage mit dem Modell
             print(f"🎯 SAM 2: Segmentiere Bereich {x1},{y1}-{x2},{y2}")
-            
-            # SAM Prediction
-            masks, scores, _ = self.sam_predictor.predict(
-                point_coords=None,
-                point_labels=None,
-                box=input_box[None, :],
-                multimask_output=False,
-                return_logits=False
-            )
-            
-            # Beste Maske extrahieren und glätten (5-Pixel Übergang)
-            mask_array = masks[0].astype(np.uint8) * 255
-            mask_array = self._smooth_mask(mask_array, blur_radius=2)  # ~5 Pixel Rand
-            
-            # Zu PIL Image konvertieren
+            with torch.no_grad():
+                outputs = self.sam_model(**inputs)
+
+            # 5. Maske extrahieren und verarbeiten
+            # outputs.pred_masks enthält die Masken-Logits
+            # post_process_masks stellt die Originalgröße wieder her
+            mask = self.sam_processor.post_process_masks(
+                outputs.pred_masks,
+                inputs.original_sizes,
+                inputs.reshaped_input_sizes
+            )[0][0]  # [batch_index][mask_index]
+
+            # Sigmoid für Wahrscheinlichkeiten, dann Schwellenwert
+            mask = mask.sigmoid().cpu().numpy()
+            mask_array = (mask > 0.5).astype(np.uint8) * 255
+
+            # 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)
+
+            # 7. 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")
-            
-            # Modus-spezifische Anpassung
+
+            # 8. Modus-spezifische Anpassung (Invertierung)
             if mode == "environment_change":
-                # MODUS 1: Umgebung ändern
-                # Objekt schwarz (0) = ERHALTEN, Umgebung weiß (255) = VERÄNDERN
+                # MODUS 1: Umgebung ändern - Objekt schwarz (erhalten)
                 mask = Image.eval(mask, lambda x: 255 - x)
                 print("   SAM-Modus: Umgebung ändern (Objekt erhalten)")
             else:
-                # MODUS 2 & 3: Focus oder Gesicht ändern
-                # Objekt weiß (255) = VERÄNDERN, Umgebung schwarz (0) = ERHALTEN
+                # MODUS 2 & 3: Focus/Gesicht ändern - Objekt weiß (verändern)
                 print("   SAM-Modus: Focus/Gesicht ändern (Objekt verändern)")
-            
+
             print(f"✅ SAM 2: Präzise Maske erstellt ({mask.size})")
             return mask
-            
+
         except Exception as e:
-            print(f"⚠️ SAM 2 Fehler: {str(e)[:100]}")
+            print(f"⚠️ SAM 2 Fehler (Transformers API): {str(e)[:200]}")
+            import traceback
+            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
-        
+
         mask = Image.new("L", image.size, 0)
-        
+
         if bbox_coords and all(coord is not None for coord in bbox_coords):
             x1, y1, x2, y2 = self._validate_bbox(image, bbox_coords)
             draw = ImageDraw.Draw(mask)
-            
+
             if mode == "environment_change":
                 # MODUS 1: Alles außer Box verändern
                 draw.rectangle([0, 0, image.size[0], image.size[1]], fill=255)
                 draw.rectangle([x1, y1, x2, y2], fill=0)
+                print("ℹ️ Rechteckige Maske: Umgebung ändern")
             else:
                 # MODUS 2 & 3: Nur Box verändern
                 draw.rectangle([x1, y1, x2, y2], fill=255)
-        
-        print("ℹ️ Rechteckige Maske (SAM Fallback)")
+                print("ℹ️ Rechteckige Maske: Focus/Gesicht ändern")
+
         return mask
-    
+
     def load_pose_detector(self):
         """Lädt nur den Pose-Detector"""
         if self.pose_detector is None:
@@ -249,35 +218,35 @@ def _lazy_load_sam(self):
             except Exception as e:
                 print(f"⚠️ Pose-Detector konnte nicht geladen werden: {e}")
         return self.pose_detector
-    
+
     def load_midas_model(self):
         """Lädt MiDaS Model für Depth Maps"""
         if self.midas_model is None:
             print("🔄 Lade MiDaS Modell für Depth Maps...")
             try:
                 import torchvision.transforms as T
-                
+
                 self.midas_model = torch.hub.load(
-                    "intel-isl/MiDaS", 
-                    "DPT_Hybrid", 
+                    "intel-isl/MiDaS",
+                    "DPT_Hybrid",
                     trust_repo=True
                 )
-                
+
                 self.midas_model.to(self.device)
                 self.midas_model.eval()
-                
+
                 self.midas_transform = T.Compose([
                     T.Resize(384),
                     T.ToTensor(),
                     T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
                 ])
-                
+
                 print("✅ MiDaS Modell erfolgreich geladen")
             except Exception as e:
                 print(f"❌ MiDaS konnte nicht geladen werden: {e}")
                 print("ℹ️ Verwende Fallback-Methode")
                 self.midas_model = None
-        
+
         return self.midas_model
 
     def extract_pose_simple(self, image):
@@ -309,13 +278,13 @@ def _lazy_load_sam(self):
         """Extrahiert Canny Edges für Umgebungserhaltung"""
         try:
             img_array = np.array(image.convert("RGB"))
-            
+
             gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
             edges = cv2.Canny(gray, 100, 200)
-            
+
             edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
             edges_image = Image.fromarray(edges_rgb)
-            
+
             print("✅ Canny Edge Map erstellt")
             return edges_image
         except Exception as e:
@@ -330,11 +299,11 @@ def _lazy_load_sam(self):
             midas = self.load_midas_model()
             if midas is not None:
                 print("🎯 Verwende MiDaS für Depth Map...")
-                
+
                 import torchvision.transforms as T
-                
+
                 img_transformed = self.midas_transform(image).unsqueeze(0).to(self.device)
-                
+
                 with torch.no_grad():
                     prediction = midas(img_transformed)
                     prediction = torch.nn.functional.interpolate(
@@ -343,22 +312,22 @@ def _lazy_load_sam(self):
                         mode="bicubic",
                         align_corners=False,
                     ).squeeze()
-                
+
                 depth_np = prediction.cpu().numpy()
                 depth_min, depth_max = depth_np.min(), depth_np.max()
-                
+
                 if depth_max > depth_min:
                     depth_np = (depth_np - depth_min) / (depth_max - depth_min)
-                
+
                 depth_np = (depth_np * 255).astype(np.uint8)
                 depth_image = Image.fromarray(depth_np).convert("RGB")
-                
+
                 print("✅ MiDaS Depth Map erfolgreich erstellt")
                 return depth_image
-            
+
             else:
                 raise Exception("MiDaS nicht geladen")
-                
+
         except Exception as e:
             print(f"⚠️ MiDaS Fehler: {e}. Verwende Fallback...")
             try:
@@ -368,7 +337,7 @@ def _lazy_load_sam(self):
                 depth_map = cv2.GaussianBlur(gray, (5, 5), 0)
                 depth_rgb = cv2.cvtColor(depth_map, cv2.COLOR_GRAY2RGB)
                 depth_image = Image.fromarray(depth_rgb)
-                
+
                 print("✅ Fallback Depth Map erstellt")
                 return depth_image
             except Exception as fallback_error:
@@ -380,7 +349,7 @@ def _lazy_load_sam(self):
         ERSTELLT NUR CONDITIONING-MAPS, generiert KEIN Bild.
         """
         print("🎯 ControlNet: Erstelle Conditioning-Maps...")
-        
+
         if keep_environment:
             print("   Modus: Depth + Canny")
             conditioning_images = [
@@ -393,7 +362,7 @@ def _lazy_load_sam(self):
                 self.extract_pose(image),
                 self.extract_canny_edges(image)
             ]
-        
+
         print(f"✅ {len(conditioning_images)} Conditioning-Maps erstellt.")
         return conditioning_images