revert: to Florence2 model

Dockerfile

@@ -29,6 +29,7 @@ RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
 
 # 6. Создаем пользователя и директории, меняем владельца
+# Это гарантирует, что все последующие операции будут иметь правильные права
 RUN useradd -m -u 1000 appuser && \
     mkdir -p /data/.cache && \
     chown -R appuser:appuser /app /data
@@ -37,13 +38,13 @@ RUN useradd -m -u 1000 appuser && \
 USER appuser
 
 # 8. Скачиваем модели от имени этого пользователя.
-# --- ИСПРАВЛЕНИЕ: Используем правильные классы для Grounding DINO ---
+# Теперь они попадут в правильный кэш, доступный приложению.
 RUN iopaint download --model lama && \
     python -c "\
-from transformers import GroundingDinoProcessor, GroundingDinoForObjectDetection; \
-model_id = 'IDEA-Research/grounding-dino-base'; \
-GroundingDinoProcessor.from_pretrained(model_id); \
-GroundingDinoForObjectDetection.from_pretrained(model_id)"
+from transformers import AutoModelForCausalLM, AutoProcessor; \
+model_id = 'microsoft/Florence-2-large'; \
+AutoProcessor.from_pretrained(model_id, trust_remote_code=True); \
+AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)"
 
 # 9. Открываем порт и запускаем приложение
 EXPOSE 7860

logic.py

@@ -1,108 +1,116 @@
 import cv2
 import numpy as np
-from PIL import Image
-from transformers import GroundingDinoProcessor, GroundingDinoForObjectDetection
+from PIL import Image, ImageDraw
+from transformers import AutoProcessor, AutoModelForCausalLM
 from iopaint.model_manager import ModelManager
 from iopaint.schema import HDStrategy, LDMSampler, InpaintRequest
 import torch
-import supervision as sv
 from loguru import logger
+from enum import Enum
 import time
 
+class TaskType(str, Enum):
+    OPEN_VOCAB_DETECTION = "<OPEN_VOCABULARY_DETECTION>"
+
 class WatermarkRemover:
     def __init__(self, device="cpu"):
         self.device = device
-        self.dino_processor = None
-        self.dino_model = None
-        self.inpainting_model = None
-        logger.info(f"WatermarkRemover инициализирован для устройства {self.device}. Модели будут загружены по первому требованию.")
-
-    def _get_dino_components(self):
-        if self.dino_model is None:
-            logger.info("Первый вызов: Загрузка модели детектора Grounding DINO...")
-            dino_model_id = "IDEA-Research/grounding-dino-base"
-            self.dino_processor = GroundingDinoProcessor.from_pretrained(dino_model_id)
-            self.dino_model = GroundingDinoForObjectDetection.from_pretrained(dino_model_id).to(self.device)
-            logger.info("Модель Grounding DINO загружена.")
-        return self.dino_processor, self.dino_model
+        logger.info(f"Используемое устройство: {self.device}")
+        
+        logger.info("Загрузка модели Florence-2...")
+        model_id = "microsoft/Florence-2-large"
+        self.florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to(self.device).eval()
+        self.florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+        logger.info("Модель Florence-2 загружена.")
 
-    def _get_inpainting_model(self):
-        if self.inpainting_model is None:
-            logger.info("Первый вызов: Загрузка модели LaMA для закрашивания...")
-            self.inpainting_model = ModelManager(name="lama", device=self.device)
-            logger.info("Модель LaMA загружена.")
-        return self.inpainting_model
+        logger.info("Загрузка модели LaMA для закрашивания...")
+        self.inpainting_model = ModelManager(name="lama", device=self.device)
+        logger.info("Модель LaMA загружена.")
 
-    def _detect_watermarks(self, image: Image, text_prompt: str, box_threshold: float, text_threshold: float):
-        processor, model = self._get_dino_components()
+    def _identify(self, image: Image.Image, text_input: str):
+        task_prompt = TaskType.OPEN_VOCAB_DETECTION
+        prompt = task_prompt.value + text_input
+        inputs = self.florence_processor(text=prompt, images=image, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
         
-        inputs = processor(images=image, text=text_prompt, return_tensors="pt").to(self.device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        results = processor.post_process_grounded_object_detection(
-            outputs,
+        # Возвращаем проверенное значение num_beams
+        generated_ids = self.florence_model.generate(
             input_ids=inputs["input_ids"],
-            box_threshold=box_threshold,
-            text_threshold=text_threshold,
-            target_sizes=[image.size[::-1]]
-        )[0]
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=1024,
+            do_sample=False,
+            num_beams=5
+        )
         
-        return sv.Detections(
-            xyxy=results["boxes"].cpu().numpy(),
-            confidence=results["scores"].cpu().numpy()
+        generated_text = self.florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        return self.florence_processor.post_process_generation(
+            generated_text, task=task_prompt.value, image_size=image.size
         )
+    
+    def _get_mask_for_prompt(self, image: Image.Image, text_input: str, max_bbox_percent: float) -> Image.Image:
+        logger.info(f"Поиск по промпту: '{text_input}'...")
+        parsed_answer = self._identify(image, text_input)
 
+        mask = Image.new("L", image.size, 0)
+        draw = ImageDraw.Draw(mask)
+        
+        detection_key = TaskType.OPEN_VOCAB_DETECTION.value
+        if detection_key in parsed_answer and "bboxes" in parsed_answer[detection_key]:
+            bboxes = parsed_answer[detection_key]["bboxes"]
+            logger.success(f"Найдено {len(bboxes)} bbox по промпту '{text_input}'.")
+            image_area = image.width * image.height
+            
+            for bbox in bboxes:
+                x1, y1, x2, y2 = map(int, bbox)
+                bbox_area = (x2 - x1) * (y2 - y1)
+                
+                if (bbox_area / image_area) * 100 <= max_bbox_percent:
+                    draw.rectangle([x1, y1, x2, y2], fill=255)
+                else:
+                    logger.warning(f"Пропущен bbox >{max_bbox_percent}%: {bbox}")
+        else:
+            logger.warning(f"Промпт '{text_input}' не дал результатов.")
+            
+        return mask
+        
     def _inpaint_image(self, image_np: np.ndarray, mask_np: np.ndarray) -> np.ndarray:
-        inpainting_model = self._get_inpainting_model()
-        config = InpaintRequest(ldm_steps=50, ldm_sampler=LDMSampler.ddim, hd_strategy=HDStrategy.CROP)
-        result = inpainting_model(image_np, mask_np, config)
+        config = InpaintRequest(ldm_steps=50, ldm_sampler=LDMSampler.ddim, hd_strategy=HDStrategy.CROP, hd_strategy_crop_margin=64, hd_strategy_crop_trigger_size=800, hd_strategy_resize_limit=1600)
+        result = self.inpainting_model(image_np, mask_np, config)
         if result.dtype in [np.float64, np.float32]:
             result = np.clip(result, 0, 255).astype(np.uint8)
         return result
 
-    def run(self, image: Image.Image) -> Image.Image:
+    def run(self, image: Image.Image, max_bbox_percent: float = 27.0) -> Image.Image:
         start_time = time.time()
         logger.info("Начало процесса удаления вотермарок...")
-
-        # --- Этап 1: Поиск высококонтрастных частей (текст) ---
-        logger.info("--- ЭТАП 1: Поиск текста с высокой уверенностью ---")
-        detections_text = self._detect_watermarks(
-            image, text_prompt="text . watermark", box_threshold=0.4, text_threshold=0.4
-        )
-        logger.info(f"Найдено {len(detections_text)} текстовых фрагментов.")
         
-        # --- Этап 2: Поиск низкоконтрастных частей (фон) ---
-        logger.info("--- ЭТАП 2: Поиск фона с низкой уверенностью ---")
-        detections_bg = self._detect_watermarks(
-            image, text_prompt="semi-transparent background . transparent overlay . watermark background", box_threshold=0.25, text_threshold=0.25
-        )
-        logger.info(f"Найдено {len(detections_bg)} фоновых фрагментов.")
+        mask_image = self._get_mask_for_prompt(image, "watermark", max_bbox_percent)
+        mask_np = np.array(mask_image)
         
-        # --- Этап 3: Объединение и очистка ---
-        # Складываем все найденные детекции вместе
-        all_detections_list = [detections_text, detections_bg]
-        detections = sv.Detections.merge(all_detections_list)
+        if not np.any(mask_np):
+            logger.warning("Основной промпт не нашел вотермарок. Запускаю резервные попытки...")
+            fallback_prompts = ["logo", "text overlay"] # Убрали слишком общие промпты
+            
+            for prompt in fallback_prompts:
+                mask_image = self._get_mask_for_prompt(image, prompt, max_bbox_percent)
+                mask_np = np.array(mask_image)
+                if np.any(mask_np):
+                    logger.success(f"Резервный промпт '{prompt}' нашел вотермарку!")
+                    break
 
-        if len(detections) == 0:
-            logger.info("Ни один из этапов не нашел вотермарок. Возвращаем оригинальное изображение.")
+        if not np.any(mask_np):
+            logger.info("Ни одна из попыток не нашла вотермарок. Возвращаем оригинальное изображение.")
             return image
-        
-        logger.info(f"Всего найдено {len(detections)} фрагментов. Запускаю NMS для слияния...")
-        # Применяем NMS для слияния всех пересекающихся рамок в одну
-        merged_detections = detections.with_nms(class_agnostic=True, threshold=0.6)
-        logger.success(f"После слияния осталось {len(merged_detections)} уникальных областей.")
 
-        # --- Этап 4: Создание маски и закрашивание ---
-        mask = np.zeros(image.size[::-1], dtype=np.uint8)
-        for box in merged_detections.xyxy:
-            cv2.rectangle(mask, tuple(map(int, box[:2])), tuple(map(int, box[2:])), 255, -1)
-            
-        kernel = np.ones((9, 9), np.uint8)
-        processed_mask = cv2.dilate(mask, kernel, iterations=1)
-            
+        kernel_size = 7
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
+        processed_mask = cv2.dilate(mask_np, kernel, iterations=1)
+        
         image_np = np.array(image.convert("RGB"))
+
+        logger.info("Закрашивание области по обработанной маске...")
         result_np_bgr = self._inpaint_image(image_np, processed_mask)
+        
         result_np_rgb = cv2.cvtColor(result_np_bgr, cv2.COLOR_BGR2RGB)
         
         end_time = time.time()

requirements.txt

@@ -4,15 +4,13 @@ uvicorn[standard]
 python-multipart
 requests
 
-# Основная ML-библиотека (устанавливает transformers, opencv и т.д.)
+# Основная ML-библиотека (она сама подтянет transformers, diffusers и т.д.)
 iopaint
 
-# --- Grounding DINO ---
-timm
-supervision
-
 # Утилиты
 loguru
+opencv-python-headless
 
-# Оптимизация
-flash-attn==2.5.8
+# Оптимизация для Transformer (устанавливаем отдельно для контроля)
+# Убедитесь, что эта версия совместима с вашим кодом и версией CUDA в Docker-образе
+flash-attn==2.5.8