Create sam_module.py

sam_module.py

@@ -0,0 +1,310 @@
+def create_sam_mask(self, image, bbox_coords, mode):
+    """
+    ERWEITERTE Funktion: Erstellt präzise Maske mit SAM 2
+    Restrukturierte Version mit klaren Blöcken pro Modus
+    """
+    try:
+        print("#" * 80)
+        print("# 🎯 STARTE SAM 2 SEGMENTIERUNG")
+        print("#" * 80)
+        print(f"📐 Eingabebild-Größe: {image.size}")
+        print(f"🎛️  Ausgewählter Modus: {mode}")
+        
+        # ============================================================
+        # VORBEREITUNG FÜR ALLE MODI
+        # ============================================================
+        original_image = image
+        
+        # 1. SAM2 laden
+        if not self.sam_initialized:
+            print("📥 SAM 2 ist noch nicht geladen, starte Lazy Loading...")
+            self._lazy_load_sam()
+
+        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)
+
+        # 2. Validiere BBox
+        x1, y1, x2, y2 = self._validate_bbox(image, bbox_coords)
+        original_bbox = (x1, y1, x2, y2)
+        print(f"📏 Original-BBox Größe: {x2-x1} × {y2-y1} px")
+        
+        # ============================================================
+        # BLOCK 1: ENVIRONMENT_CHANGE
+        # ============================================================
+        if mode == "environment_change":
+            print("-" * 60)
+            print("🌳 MODUS: ENVIRONMENT_CHANGE")
+            print("-" * 60)
+            
+            # ... existierende environment_change Logik hier komplett ...
+            # (wird aus dem Original übernommen, nicht verändert)
+            
+            # WICHTIG: Du musst den environment_change Code hier einfügen
+            # von Zeile ~175 bis ~250 aus dem Original
+            
+            # Beispiel-Struktur (vereinfacht):
+            image_np = np.array(image.convert("RGB"))
+            input_boxes = [[[x1, y1, x2, y2]]]
+            
+            # KEINE Punkte für environment_change
+            inputs = self.sam_processor(
+                image_np, 
+                input_boxes=input_boxes,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            with torch.no_grad():
+                outputs = self.sam_model(**inputs)
+            
+            # Nur beste Maske verwenden und auf 512x512 skalieren
+            best_mask = outputs.pred_masks[:, :, 0, :, :]  # Erste Maske nehmen
+            resized_mask = F.interpolate(
+                best_mask,
+                size=(512, 512),  # Direkt auf ControlNet-Zielgröße
+                mode='bilinear',
+                align_corners=False
+            ).squeeze()
+            
+            mask_np = resized_mask.sigmoid().cpu().numpy()
+            
+            # Invertieren für environment_change
+            threshold = 0.5
+            mask_array = (mask_np > threshold).astype(np.uint8) * 255
+            mask_array = 255 - mask_array  # Invertieren
+            
+            # Auf Originalgröße für Rückgabe
+            mask = Image.fromarray(mask_array).convert("L")
+            mask = mask.resize(original_image.size, Image.Resampling.NEAREST)
+            
+            return mask, mask  # raw_mask gleiche wie finale Maske
+            
+        # ============================================================
+        # BLOCK 2: FOCUS_CHANGE (KORRIGIERTE VERSION)
+        # ============================================================
+        elif mode == "focus_change":
+            print("-" * 60)
+            print("🎯 MODUS: FOCUS_CHANGE (OPTIMIERT)")
+            print("-" * 60)
+            
+            # Bild für SAM vorbereiten
+            image_np = np.array(image.convert("RGB"))
+            
+            # NUR EINE BBOX UND NUR MITTELPUNKT (kein Gesichtspunkt)
+            input_boxes = [[[x1, y1, x2, y2]]]
+            
+            # Nur Mittelpunkt als positiver Prompt
+            center_x = (x1 + x2) // 2
+            center_y = (y1 + y2) // 2
+            input_points = [[[[center_x, center_y]]]]  # NUR EIN PUNKT
+            input_labels = [[[1]]]  # Positiver Prompt
+            
+            print(f"   🎯 SAM-Prompt: BBox [{x1},{y1},{x2},{y2}]")
+            print(f"   👁️  Punkt: Nur Mitte ({center_x},{center_y})")
+            
+            # SAM Inputs vorbereiten
+            inputs = self.sam_processor(
+                image_np, 
+                input_boxes=input_boxes,
+                input_points=input_points,
+                input_labels=input_labels,
+                return_tensors="pt"
+            ).to(self.device)
+            
+            # SAM Vorhersage (alle 3 Masken)
+            print("🧠 SAM 2 INFERENZ (3 Masken-Varianten)")
+            with torch.no_grad():
+                outputs = self.sam_model(**inputs)
+            
+            # BBox-Information für Heuristik
+            bbox_center = ((x1 + x2) // 2, (y1 + y2) // 2)
+            bbox_area = (x2 - x1) * (y2 - y1)
+            
+            print("🤔 HEURISTIK: Beste Maske auswählen")
+            best_mask_idx = 0
+            best_score = -1
+            
+            # Alle 3 Masken analysieren (OHNE sie alle zu skalieren!)
+            for i in range(3):
+                # Maske in Original-SAM-Größe (256x256) analysieren
+                mask_256 = outputs.pred_masks[:, :, i, :, :]
+                mask_np_256 = mask_256.sigmoid().squeeze().cpu().numpy()
+                
+                # Für Heuristik: Temporär auf Bildgröße skalieren
+                temp_mask = F.interpolate(
+                    mask_256,
+                    size=(image.height, image.width),
+                    mode='bilinear',
+                    align_corners=False
+                ).squeeze()
+                mask_np_temp = temp_mask.sigmoid().cpu().numpy()
+                
+                # Adaptive Vor-Filterung
+                mask_max = mask_np_temp.max()
+                if mask_max < 0.3:
+                    continue  # Maske überspringen
+                
+                adaptive_threshold = max(0.3, mask_max * 0.7)
+                mask_binary = (mask_np_temp > adaptive_threshold).astype(np.uint8)
+                
+                if np.sum(mask_binary) == 0:
+                    continue
+                
+                # Heuristik-Berechnung (wie bisher)
+                mask_area_pixels = np.sum(mask_binary)
+                
+                # BBox-Überlappung
+                bbox_mask = np.zeros((image.height, image.width), dtype=np.uint8)
+                bbox_mask[y1:y2, x1:x2] = 1
+                overlap = np.sum(mask_binary & bbox_mask)
+                bbox_overlap_ratio = overlap / np.sum(bbox_mask) if np.sum(bbox_mask) > 0 else 0
+                
+                # Schwerpunkt
+                y_coords, x_coords = np.where(mask_binary > 0)
+                if len(y_coords) > 0:
+                    centroid_y = np.mean(y_coords)
+                    centroid_x = np.mean(x_coords)
+                    centroid_distance = np.sqrt((centroid_x - bbox_center[0])**2 + 
+                                                (centroid_y - bbox_center[1])**2)
+                    normalized_distance = centroid_distance / max(image.width, image.height)
+                else:
+                    normalized_distance = 1.0
+                
+                # Flächen-Ratio
+                area_ratio = mask_area_pixels / bbox_area
+                area_score = 1.0 - min(abs(area_ratio - 1.0), 1.0)
+                
+                # FOCUS_CHANGE spezifischer Score
+                score = (
+                    bbox_overlap_ratio * 0.4 +      # 40% BBox-Überlappung
+                    (1.0 - normalized_distance) * 0.25 +  # 25% Zentrumsnähe
+                    area_score * 0.25 +             # 25% Flächenpassung
+                    mask_max * 0.1                  # 10% SAM-Konfidenz
+                )
+                
+                print(f"   Maske {i+1}: Score={score:.3f}, "
+                      f"Überlappung={bbox_overlap_ratio:.3f}, "
+                      f"Fläche={mask_area_pixels:,}px")
+                
+                if score > best_score:
+                    best_score = score
+                    best_mask_idx = i
+            
+            print(f"✅ Beste Maske: Nr. {best_mask_idx+1} mit Score {best_score:.3f}")
+            
+            # NUR DIE BESTE MASKE AUF 512x512 SKALIEREN
+            best_mask_256 = outputs.pred_masks[:, :, best_mask_idx, :, :]
+            resized_mask = F.interpolate(
+                best_mask_256,
+                size=(512, 512),  # DIREKT AUF CONTROLNET-ZIELGRÖßE
+                mode='bilinear',
+                align_corners=False
+            ).squeeze()
+            
+            mask_np = resized_mask.sigmoid().cpu().numpy()
+            print(f"   🔄 Beste Maske skaliert auf 512×512 für ControlNet")
+            
+            # Dynamischer Threshold für focus_change
+            mask_max = mask_np.max()
+            if best_score < 0.7:  # Schlechte Maskenqualität
+                dynamic_threshold = 0.05  # SEHR NIEDRIG für maximale Abdeckung
+                print(f"   ⚠️  Masken-Score niedrig ({best_score:.3f}). "
+                      f"Threshold=0.05 für maximale Abdeckung")
+            else:
+                dynamic_threshold = max(0.15, mask_max * 0.3)  # Moderater Threshold
+                print(f"   ✅ Gute Maske. Threshold={dynamic_threshold:.3f}")
+            
+            # Binärmaske erstellen
+            mask_array = (mask_np > dynamic_threshold).astype(np.uint8) * 255
+            
+            # Fallback bei leerer Maske
+            if mask_array.max() == 0:
+                print("   ⚠️  Maske leer, erstelle rechteckige Fallback-Maske")
+                mask_array = np.zeros((512, 512), dtype=np.uint8)
+                # BBox auf 512x512 skalieren für Fallback
+                scale_x = 512 / image.width
+                scale_y = 512 / image.height
+                fb_x1 = int(x1 * scale_x)
+                fb_y1 = int(y1 * scale_y)
+                fb_x2 = int(x2 * scale_x)
+                fb_y2 = int(y2 * scale_y)
+                cv2.rectangle(mask_array, (fb_x1, fb_y1), (fb_x2, fb_y2), 255, -1)
+            
+            # FOCUS_CHANGE POSTPROCESSING (angepasst für 512x512)
+            print("🔧 FOCUS_CHANGE POSTPROCESSING (auf 512×512)")
+            
+            # 1. Größte Komponente behalten
+            labeled_array, num_features = ndimage.label(mask_array)
+            if num_features > 1:
+                sizes = ndimage.sum(mask_array, labeled_array, range(1, num_features + 1))
+                largest_component = np.argmax(sizes) + 1
+                mask_array = np.where(labeled_array == largest_component, mask_array, 0)
+                print(f"   ✅ Größte Komponente behalten ({num_features}→1)")
+            
+            # 2. Morphologische Operationen
+            kernel_close = np.ones((5, 5), np.uint8)
+            mask_array = cv2.morphologyEx(mask_array, cv2.MORPH_CLOSE, kernel_close, iterations=2)
+            
+            kernel_dilate = np.ones((15, 15), np.uint8)
+            mask_array = cv2.dilate(mask_array, kernel_dilate, iterations=1)
+            
+            # 3. Weiche Übergänge
+            mask_array = cv2.GaussianBlur(mask_array, (9, 9), 2.0)
+            
+            # 4. Gamma-Korrektur
+            mask_array_float = mask_array.astype(np.float32) / 255.0
+            mask_array_float = np.clip(mask_array_float, 0.0, 1.0)
+            mask_array_float = mask_array_float ** 0.85
+            mask_array = (mask_array_float * 255).astype(np.uint8)
+            
+            # 5. Auf Originalgröße für Rückgabe (falls benötigt)
+            mask_512 = Image.fromarray(mask_array).convert("L")
+            raw_mask = mask_512.copy()  # Rohmaske = finale Maske bei focus_change
+            
+            # Finale Maske für ControlNet ist 512x512
+            mask = mask_512
+            
+            print(f"✅ FOCUS_CHANGE Maske erstellt: {mask.size}")
+            return mask, raw_mask
+            
+        # ============================================================
+        # BLOCK 3: FACE_ONLY_CHANGE
+        # ============================================================
+        elif mode == "face_only_change":
+            print("-" * 60)
+            print("👤 MODUS: FACE_ONLY_CHANGE")
+            print("-" * 60)
+            
+            # ... existierende face_only_change Logik hier komplett ...
+            # (wird aus dem Original übernommen, nicht verändert)
+            
+            # WICHTIG: Du musst den face_only_change Code hier einfügen
+            # von Zeile ~252 bis ~650 aus dem Original
+            
+            # Beispiel-Struktur (vereinfacht):
+            # Crop, Punkte setzen, spezielle Gesichtsheuristik etc.
+            
+            # Am Ende:
+            mask = Image.new("L", (512, 512), 128)  # Platzhalter
+            raw_mask = mask.copy()
+            return mask, raw_mask
+            
+        # ============================================================
+        # UNBEKANNTER MODUS
+        # ============================================================
+        else:
+            print(f"❌ Unbekannter Modus: {mode}")
+            return self._create_rectangular_mask(image, bbox_coords, "focus_change")
+            
+    except Exception as e:
+        print("❌ FEHLER IN SAM 2 SEGMENTIERUNG")
+        print(f"Fehler: {str(e)[:200]}")
+        import traceback
+        traceback.print_exc()
+        
+        # Fallback
+        fallback_mask = self._create_rectangular_mask(original_image, original_bbox, mode)
+        if fallback_mask.size != original_image.size:
+            fallback_mask = fallback_mask.resize(original_image.size, Image.Resampling.NEAREST)
+        
+        return fallback_mask, fallback_mask