Update app.py

app.py

@@ -79,12 +79,7 @@ class FreeAIOrchestrator:
             print("🔄 Инициализация кастомного токенайзера для Surya Table...")
 
             # Используем пайплайн с указанием нашего токенайзера
-            self.surya_table_model = pipeline(
-                "image-to-text",
-                model="vikp/surya_table",
-                tokenizer=Byt5LangTokenizer.from_pretrained("t5-base", legacy=False)
-            )
-
+            self.surya_table_model = TableRecPredictor()
             print("✅ Surya Table модель загружена успешно")
             self.surya_table_available = True
         except Exception as e:
@@ -255,62 +250,25 @@ class FreeAIOrchestrator:
                         pil_image = Image.open(io.BytesIO(image))
 
                     # Распознаем таблицу через Surya Table
-                    try:
-                        # Пробуем с нашим кастомным токенайзером
-                        table_result = self.surya_table_model(pil_image)
-                    except Exception as tokenizer_error:
-                        print(f"⚠️ Ошибка с кастомным токенайзером: {tokenizer_error}")
-                        # Если не сработало, используем альтернативный метод
-                        try:
-                            print("🔄 Используем альтернативный метод для Surya Table...")
-                            # Используем автоматически выбранный tokenizer
-                            from transformers import pipeline
-                            surya_fallback = pipeline(
-                                "image-to-text",
-                                model="vikp/surya_table",
-                                tokenizer=None  # Позволяем pipeline самому выбрать tokenizer
-                            )
-                            table_result = surya_fallback(pil_image)
-                        except Exception as fallback_error:
-                            print(f"⚠️ Ошибка с fallback методом: {fallback_error}")
-                            # Если и это не сработало, используем еще более простой метод
-                            try:
-                                # Альтернативный подход без использования pipeline
-                                print("🔄 Используем прямой подход для Surya Table...")
-                                from transformers import AutoModelForVision2Seq, AutoImageProcessor
-
-                                processor = AutoImageProcessor.from_pretrained("vikp/surya_table")
-                                model = AutoModelForVision2Seq.from_pretrained("vikp/surya_table")
-
-                                inputs = processor(images=pil_image, return_tensors="pt")
-                                outputs = model.generate(**inputs, max_length=1024)
-
-                                table_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-                                table_result = [{"generated_text": table_text}]
-                            except Exception as direct_error:
-                                print(f"⚠️ Все методы распознавания Surya Table не удались: {direct_error}")
-                                raise direct_error
+                    table_result = self.surya_table_model([pil_image])
 
                     # Преобразуем результат в структурированный формат
                     try:
                         # Результат может быть в разных форматах
                         if isinstance(table_result, list) and len(table_result) > 0:
-                            if isinstance(table_result[0], dict) and 'generated_text' in table_result[0]:
-                                table_text = table_result[0]['generated_text']
+                            if isinstance(table_result[0], dict) and 'cells' in table_result[0]:
+                                table_data = self._parse_surya_table(table_result[0])
                             else:
-                                table_text = str(table_result)
+                                table_data = []
                         else:
-                            table_text = str(table_result)
-
-                        # Парсим структуру таблицы из текста
-                        table_data = self._parse_table_text(table_text)
+                            table_data = []
 
                         return {
                             "success": True,
                             "type": "table",
                             "model": "surya_table",
                             "rows": table_data,
-                            "raw_text": table_text,
+                            "raw_text": "",
                             "confidence": 0.95
                         }
                     except Exception as parse_error:
@@ -347,6 +305,55 @@ class FreeAIOrchestrator:
             print(f"❌ Ошибка распознавания таблицы: {e}")
             return []
 
+    def _parse_surya_table(self, surya_result):
+        """Парсинг результата Surya в структурированные данные"""
+        rows = []
+        headers = {}
+
+        # First pass: get headers
+        for cell in surya_result.get('cells', []):
+            if cell.get('is_header'):
+                headers[cell.get('col_id')] = cell.get('text', '').lower()
+
+        # Second pass: get rows
+        row_dict = {}
+        for cell in surya_result.get('cells', []):
+            row_id = cell.get('row_id')
+            if row_id not in row_dict:
+                row_dict[row_id] = {}
+
+            col_id = cell.get('col_id')
+            header = headers.get(col_id, str(col_id))
+
+            value = cell.get('text', '')
+            # Преобразуем заголовки к стандартным полям
+            if 'товар' in header or 'название' in header or 'наимен' in header:
+                row_dict[row_id]['name'] = value
+            elif 'кол' in header or 'шт' in header:
+                try:
+                    # Извлекаем числовое значение
+                    quantity = re.search(r'(\d+(?:\.\d+)?)', value)
+                    if quantity:
+                        row_dict[row_id]['quantity'] = float(quantity.group(1))
+                    else:
+                        row_dict[row_id]['quantity'] = value
+                except:
+                    row_dict[row_id]['quantity'] = value
+            elif 'арт' in header:
+                row_dict[row_id]['article'] = value
+            elif 'цен' in header:
+                # Извлекаем числовое значение цены
+                price = re.search(r'(\d+(?:\.\d+)?)', value)
+                if price:
+                    row_dict[row_id]['price'] = float(price.group(1))
+                else:
+                    row_dict[row_id]['price'] = value
+            else:
+                # Для прочих колонок используем оригинальное название
+                row_dict[row_id][header] = value
+
+        return list(row_dict.values())
+
     def _parse_table_text(self, table_text):
         """Парсинг текста таблицы в структурированные данные"""
         rows = []