Update app.py

app.py

@@ -53,33 +53,67 @@ class FreeAIOrchestrator:
     def __init__(self):
         print("🚀 Инициализация SkladBot Free AI...")
 
-        # TrOCR для печатного текста (БЕСПЛАТНО)
-        self.printed_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
-        self.printed_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
-
-        # TrOCR для рукописного текста (БЕСПЛАТНО)
-        self.handwritten_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-        self.handwritten_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
-
-        # LayoutLM для понимания документов (БЕСПЛАТНО)
-        self.document_qa = pipeline(
-            "document-question-answering",
-            model="impira/layoutlm-document-qa"
-        )
+        # Для предотвращения ошибок при запуске в HF Space
+        try:
+            # TrOCR для печатного текста (БЕСПЛАТНО)
+            self.printed_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
+            self.printed_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
+            print("✅ TrOCR (печатный) загружен успешно")
+        except Exception as e:
+            print(f"⚠️ Ошибка при загрузке TrOCR печатный: {e}")
+            self.printed_processor = None
+            self.printed_model = None
 
-        # Table Transformer для таблиц (БЕСПЛАТНО)
-        self.table_detector = pipeline(
-            "object-detection",
-            model="microsoft/table-transformer-structure-recognition"
-        )
+        try:
+            # TrOCR для рукописного текста (БЕСПЛАТНО)
+            self.handwritten_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+            self.handwritten_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+            print("✅ TrOCR (рукописный) загружен успешно")
+        except Exception as e:
+            print(f"⚠️ Ошибка при загрузке TrOCR рукописный: {e}")
+            self.handwritten_processor = None
+            self.handwritten_model = None
+
+        try:
+            # LayoutLM для понимания документов (БЕСПЛАТНО)
+            self.document_qa = pipeline(
+                "document-question-answering",
+                model="impira/layoutlm-document-qa"
+            )
+            print("✅ LayoutLM загружен успешно")
+        except Exception as e:
+            print(f"⚠️ Ошибка при загрузке LayoutLM: {e}")
+            self.document_qa = None
+
+        try:
+            # Table Transformer для таблиц (БЕСПЛАТНО)
+            self.table_detector = pipeline(
+                "object-detection",
+                model="microsoft/table-transformer-structure-recognition"
+            )
+            print("✅ Table Transformer загружен успешно")
+        except Exception as e:
+            print(f"⚠️ Ошибка при загрузке Table Transformer: {e}")
+            self.table_detector = None
+
+        # Проинициализируем переменную для NER
+        self.ner_pipeline = None
+        try:
+            # Инициализируем NER модель для русского языка
+            self.ner_pipeline = pipeline(
+                "token-classification",
+                model="Gherman/bert-base-NER-Russian"
+            )
+            print("✅ Russian NER загружен успешно")
+        except Exception as e:
+            print(f"⚠️ Ошибка при загрузке NER: {e}")
 
-        # NEW: Добавляем интеграцию с Surya Table (БЕСПЛАТНО)
+        # Пробуем инициализировать Surya Table модель
+        self.surya_table_available = False
         try:
-            # Регистрируем кастомный токенайзер перед загрузкой модели
-            print("🔄 Инициализация кастомного токенайзера для Surya Table...")
+            print("🔄 Инициализация модели Surya Table...")
 
-            # Используем пайплайн с указанием стандартного токенайзера вместо кастомного
-            # Это решает проблему совместимости
+            # Используем стандартный токенайзер вместо кастомного
             self.surya_table_model = pipeline(
                 "image-to-text",
                 model="vikp/surya_tablerec",
@@ -90,7 +124,7 @@ class FreeAIOrchestrator:
             self.surya_table_available = True
         except Exception as e:
             print(f"⚠️ Не удалось загрузить Surya Table: {e}")
-            self.surya_table_available = False
+            self.surya_table_model = None
 
         self.stats = {
             "total_requests": 0,
@@ -192,6 +226,11 @@ class FreeAIOrchestrator:
     async def extract_printed_text(self, image):
         """Извлечение печатного текста через TrOCR"""
         try:
+            # Проверяем, инициализированы ли необходимые модели
+            if self.printed_processor is None or self.printed_model is None:
+                print("⚠️ TrOCR для печатного текста не инициализирован")
+                return ""
+
             pixel_values = self.printed_processor(image, return_tensors="pt").pixel_values
             generated_ids = self.printed_model.generate(pixel_values)
             generated_text = self.printed_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
@@ -203,6 +242,11 @@ class FreeAIOrchestrator:
     async def extract_handwritten_text(self, image):
         """Извлечение рукописного текста через TrOCR"""
         try:
+            # Проверяем, инициализированы ли необходимые модели
+            if self.handwritten_processor is None or self.handwritten_model is None:
+                print("⚠️ TrOCR для рукописного текста не инициализирован")
+                return ""
+
             pixel_values = self.handwritten_processor(image, return_tensors="pt").pixel_values
             generated_ids = self.handwritten_model.generate(pixel_values)
             generated_text = self.handwritten_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
@@ -214,6 +258,11 @@ class FreeAIOrchestrator:
     async def extract_structured_data(self, image, doc_type):
         """Структурированное понимание документа через LayoutLM"""
         try:
+            # Проверяем, инициализирована ли модель
+            if self.document_qa is None:
+                print("⚠️ LayoutLM не инициализирован")
+                return {}
+
             # Определяем вопросы на основе типа документа
             questions = self.get_document_questions(doc_type)
 
@@ -222,7 +271,8 @@ class FreeAIOrchestrator:
                 try:
                     result = self.document_qa(image=image, question=question)
                     results[question] = result["answer"]
-                except:
+                except Exception as inner_e:
+                    print(f"⚠️ Ошибка запроса к LayoutLM: {inner_e}")
                     results[question] = ""
 
             return results
@@ -230,11 +280,43 @@ class FreeAIOrchestrator:
             print(f"❌ Ошибка LayoutLM: {e}")
             return {}
 
+    def get_document_questions(self, doc_type):
+        """Формирует набор вопросов для структурированного извлечения через LayoutLM"""
+        # Базовые вопросы для всех типов документов
+        base_questions = [
+            "Что это за документ?",
+            "Какие товары указаны в документе?",
+            "Какое количество товаров указано?"
+        ]
+
+        # Специфичные вопросы в зависимости от типа документа
+        if doc_type == "invoice":
+            return base_questions + [
+                "Какая общая сумма?",
+                "Кто поставщик?",
+                "Какая дата документа?",
+                "Какой номер документа?"
+            ]
+        elif doc_type == "table":
+            return base_questions + [
+                "Сколько строк в таблице?",
+                "Какие колонки есть в таблице?",
+                "Есть ли в таблице артикулы товаров?"
+            ]
+        elif doc_type == "form":
+            return base_questions + [
+                "Кто заполнил форму?",
+                "Какой статус документа?",
+                "Требуется ли подпись?"
+            ]
+        else:
+            return base_questions
+
     async def extract_table_data(self, image):
         """Извлечение табличных данных через специализированные модели"""
         try:
             # Проверка наличия модели Surya Table
-            if hasattr(self, 'surya_table_available') and self.surya_table_available:
+            if hasattr(self, 'surya_table_available') and self.surya_table_available and self.surya_table_model is not None:
                 try:
                     # Попытка использования Surya Table для структурированного распознавания таблиц
                     print("🔍 Используем Surya Table для структурированного распознавания таблицы...")
@@ -435,8 +517,11 @@ class FreeAIOrchestrator:
 
         # Используем NER для извлечения сущностей
         try:
-            entities = self.ner_pipeline(text)
-        except:
+            entities = []
+            if self.ner_pipeline is not None:
+                entities = self.ner_pipeline(text)
+        except Exception as e:
+            print(f"⚠️ Ошибка при использовании NER: {e}")
             entities = []
 
         # Регулярные выражения для складских данных
@@ -539,13 +624,107 @@ class FreeAIOrchestrator:
             return 0.0
         return round(self.stats["successful_extractions"] / self.stats["total_requests"] * 100, 1)
 
+    async def classify_document_type(self, image):
+        """Определяет тип документа на основе визуальных характеристик"""
+        try:
+            # Преобразуем в numpy массив для анализа
+            if not isinstance(image, np.ndarray):
+                if hasattr(image, 'convert'):
+                    image_np = np.array(image.convert('RGB'))
+                else:
+                    # Если не можем получить numpy массив, возвращаем значение по умолчанию
+                    return "auto"
+            else:
+                image_np = image
+
+            # Проверяем размеры и соотношение сторон
+            height, width = image_np.shape[:2]
+            aspect_ratio = width / height
+
+            # Анализируем изображение для определения типа документа
+
+            # 1. Проверка на таблицу: регулярная структура, сетка, преобладание линий
+            # Упрощенная эвристика: таблицы часто имеют большое кол-во горизонтальных/вертикальных линий
+            try:
+                # Преобразование в оттенки серого
+                if len(image_np.shape) == 3:
+                    gray = np.mean(image_np, axis=2).astype(np.uint8)
+                else:
+                    gray = image_np
+
+                # Определение градиентов для обнаружения линий
+                grad_x = np.abs(np.diff(gray, axis=1)).sum()
+                grad_y = np.abs(np.diff(gray, axis=0)).sum()
+
+                # Нормализация по размеру изображения
+                grad_x_norm = grad_x / (width * height)
+                grad_y_norm = grad_y / (width * height)
+
+                # Если много градиентов (линий) в обоих направлениях - вероятно, это таблица
+                if grad_x_norm > 0.1 and grad_y_norm > 0.1:
+                    return "table"
+            except Exception as e:
+                print(f"⚠️ Ошибка при анализе таблиц: {e}")
+
+            # 2. Проверка на счет/накладную: обычно содержит адреса, суммы и реквизиты
+            # Здесь упрощенно используем соотношение сторон: счета часто в портретной ориентации
+            if aspect_ratio < 0.9:  # Портретная ориентация
+                return "invoice"
+
+            # 3. Проверка на форму: структурированный документ с полями для заполнения
+            # Упрощенная эвристика: много белых областей с текстом по краям
+            try:
+                # Анализ распределения пикселей
+                if len(image_np.shape) == 3:
+                    # Для цветных изображений
+                    lightness = np.mean(image_np, axis=2)
+                    light_pixels = (lightness > 200).sum()
+                    light_ratio = light_pixels / (width * height)
+
+                    if light_ratio > 0.7:  # Много светлых областей
+                        return "form"
+                else:
+                    # Для ч/б изображений
+                    light_pixels = (image_np > 200).sum()
+                    light_ratio = light_pixels / (width * height)
+
+                    if light_ratio > 0.7:
+                        return "form"
+            except Exception as e:
+                print(f"⚠️ Ошибка при анализе формы: {e}")
+
+            # 4. Проверка на рукописный текст: более нерегулярная структура
+            # Рукописный текст сложно определить только по изображению без ML
+            # Это заглушка для демонстрации, в реальности нужно использовать ML классификатор
+
+            # По умолчанию считаем документ печатным
+            return "auto"
+
+        except Exception as e:
+            print(f"❌ Ошибка при определении типа документа: {e}")
+            return "auto"  # По умолчанию для безопасности
+
 # Инициализация AI
 ai_orchestrator = FreeAIOrchestrator()
 
-# Gradio интерфейс
+# Создаем глобальный экземпляр AI Orchestrator
+ai_orchestrator = None
+
+# Отложенная загрузка моделей
+def load_ai_models():
+    global ai_orchestrator
+    ai_orchestrator = FreeAIOrchestrator()
+
+# Обработчик для Gradio
 def process_warehouse_document(image, document_type):
     """Обработка складского документа через Gradio"""
     try:
+        # Инициализируем AI если еще не инициализирован
+        global ai_orchestrator
+        if ai_orchestrator is None:
+            print("🔄 Инициализация AI при первом запросе...")
+            ai_orchestrator = FreeAIOrchestrator()
+
         import asyncio
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
@@ -565,56 +744,76 @@ def process_warehouse_document(image, document_type):
 
 def get_service_stats():
     """Получение статистики сервиса"""
+    global ai_orchestrator
+    if ai_orchestrator is None:
+        return json.dumps({
+            "status": "AI модели еще не загружены",
+            "uptime": "0 часов",
+            "ready": False
+        }, ensure_ascii=False, indent=2)
+
     stats = ai_orchestrator.get_stats()
     return json.dumps(stats, ensure_ascii=False, indent=2)
 
-# Gradio интерфейс
-with gr.Blocks(title="SkladBot Free AI") as app:
-    gr.Markdown("# 🤖 SkladBot Free AI Microservice")
-    gr.Markdown("**БЕСПЛАТНАЯ** обработка складских документов через AI")
-
-    with gr.Tab("Обработка документов"):
-        image_input = gr.Image(type="pil", label="Загрузите изображение документа")
-        doc_type = gr.Dropdown(
-            choices=["auto", "invoice", "table", "form", "handwritten"],
-            value="auto",
-            label="Тип документа"
-        )
+# Создаем интерфейс
+def create_interface():
+    with gr.Blocks(title="SkladBot Free AI", theme=gr.themes.Default()) as app:
+        gr.Markdown("# 🤖 SkladBot Free AI Microservice")
+        gr.Markdown("**БЕСПЛАТНАЯ** обработка складских документов через AI")
+
+        with gr.Tab("Обработка документов"):
+            image_input = gr.Image(type="pil", label="Загрузите изображение документа")
+            doc_type = gr.Dropdown(
+                choices=["auto", "invoice", "table", "form", "handwritten"],
+                value="auto",
+                label="Тип документа"
+            )
 
-        process_btn = gr.Button("🔍 Обработать документ", variant="primary")
+            process_btn = gr.Button("🔍 Обработать документ", variant="primary")
 
-        result_output = gr.Textbox(
-            label="Результат обработки",
-            lines=20,
-            max_lines=30
-        )
+            result_output = gr.Textbox(
+                label="Результат обработки",
+                lines=20,
+                max_lines=30
+            )
 
-        process_btn.click(
-            process_warehouse_document,
-            inputs=[image_input, doc_type],
-            outputs=result_output
-        )
+            process_btn.click(
+                process_warehouse_document,
+                inputs=[image_input, doc_type],
+                outputs=result_output
+            )
 
-    with gr.Tab("Статистика"):
-        stats_btn = gr.Button("📊 Обновить статистику")
-        stats_output = gr.Textbox(
-            label="Статистика сервиса",
-            lines=10
-        )
+        with gr.Tab("Статистика"):
+            stats_btn = gr.Button("📊 Обновить статистику")
+            stats_output = gr.Textbox(
+                label="Статистика сервиса",
+                lines=10
+            )
 
-        stats_btn.click(
-            get_service_stats,
-            outputs=stats_output
-        )
+            stats_btn.click(
+                get_service_stats,
+                outputs=stats_output
+            )
+
+        gr.Markdown("---")
+        gr.Markdown("💰 **Стоимость**: $0 (100% бесплатно)")
+        gr.Markdown("📊 **Лимит**: 20,000 запросов/месяц")
+        gr.Markdown("🧠 **AI модели**: TrOCR, LayoutLM, Table Transformer, RuBERT, SuryaTable")
 
-    gr.Markdown("---")
-    gr.Markdown("💰 **Стоимость**: $0 (100% бесплатно)")
-    gr.Markdown("📊 **Лимит**: 20,000 запросов/месяц")
-    gr.Markdown("🧠 **AI модели**: TrOCR, LayoutLM, Table Transformer, RuBERT, SuryaTable")
+    return app
 
+# Создаем интерфейс
+app = create_interface()
+
+# Запуск для Hugging Face Spaces
 if __name__ == "__main__":
+    # Запускаем интерфейс Gradio сначала, затем инициализируем модели в фоне
+    import threading
+    threading.Thread(target=load_ai_models).start()
+
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
+        debug=True,
         show_error=True
     )