Update Main.py

Main.py

@@ -5,7 +5,6 @@ import torch
 from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
 from PIL import Image
 
-# Ускоряем скачивание
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 st.set_page_config(
@@ -30,7 +29,7 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025...\n(2–6 минут при первом запуске)")
+@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025... (2–6 минут при первом запуске)")
 def load_model():
     model_name = "lightonai/LightOnOCR-1B-1025"
     
@@ -61,7 +60,7 @@ def load_image():
         return Image.open(io.BytesIO(image_data)).convert('RGB')
     return None
 
-# ==================== Интерфейс ====================
+# ==================== Главный интерфейс ====================
 st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
 st.title("LightOnOCR")
 st.markdown("**Распознавание текста с изображений**")
@@ -79,20 +78,19 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
     if img is None:
         st.error("Сначала загрузите изображение")
     else:
-        with st.spinner("Распознавание текста... (5–20 сек на CPU)"):
+        with st.spinner("Распознавание текста... (может занять 5–30 сек на CPU)"):
             
-            # Шаблон разговора
             conversation = [
                 {
                     "role": "user",
                     "content": [
                         {"type": "image"},
-                        {"type": "text", "text": "Extract all the text from this image accurately. Preserve original formatting, tables, and line breaks as much as possible."}
+                        {"type": "text", "text": "Extract ALL visible text from the image as accurately as possible. Include every word, number, and line. Preserve formatting and tables."}
                     ]
                 }
             ]
 
-            # Применяем шаблон чата (без изображения)
+            # Шаблон чата
             inputs = processor.apply_chat_template(
                 conversation,
                 add_generation_prompt=True,
@@ -104,19 +102,16 @@ if st.button("🔍 Распознать текст", use_container_width=True, t
             # Обработка изображения
             image_inputs = processor.image_processor(img, return_tensors="pt")
             pixel_values = image_inputs.pixel_values.to(device=device, dtype=dtype)
-            
-            # Получаем размеры изображения (важно для этой модели!)
-            height, width = img.height, img.width
-            image_sizes = torch.tensor([[height, width]], dtype=torch.long).to(device)
 
-            # Добавляем в inputs
+            # Размеры изображения (критично для модели)
+            height, width = img.size[1], img.size[0]   # PIL: (width, height) → height, width
+            image_sizes = torch.tensor([[height, width]], dtype=torch.long, device=device)
+
             inputs["pixel_values"] = pixel_values
             inputs["image_sizes"] = image_sizes
 
-            # Переносим остальные тензоры
-
-
-for k, v in list(inputs.items()):
+            # Перенос остальных тензоров
+            for k, v in inputs.items():
                 if isinstance(v, torch.Tensor) and k not in ["pixel_values", "image_sizes"]:
                     inputs[k] = v.to(device=device)
 
@@ -131,29 +126,35 @@ for k, v in list(inputs.items()):
                 eos_token_id=processor.tokenizer.eos_token_id,
             )
 
-            # Убираем промпт
+            # Извлекаем только сгенерированную часть
             prompt_length = inputs["input_ids"].shape[1]
             generated_ids = output_ids[0, prompt_length:]
-            
+
             generated_text = processor.decode(
                 generated_ids,
                 skip_special_tokens=True,
                 clean_up_tokenization_spaces=True
             ).strip()
 
-            # Результат
-            st.success("✅ Распознавание завершено!")
-            st.markdown('<div class="result-box">', unsafe_allow_html=True)
-            st.subheader("📝 Распознанный текст")
-            st.code(generated_text, language=None)
-            st.markdown('</div>', unsafe_allow_html=True)
-
-            st.download_button(
-                label="💾 Скачать как .txt",
-                data=generated_text,
-                file_name="recognized_text.txt",
-                mime="text/plain"
-            )
+            # Отладка (временно п��казываем длину)
+            st.info(f"Сгенерировано токенов: {len(generated_ids)} | Длина текста: {len(generated_text)} символов")
+
+            if not generated_text or len(generated_text) < 5:
+                st.warning("Модель вернула очень короткий или пустой текст. Попробуйте другое изображение с чётким английским текстом.")
+                st.code("Результат пустой или слишком короткий.", language=None)
+            else:
+                st.success("✅ Распознавание завершено!")
+                st.markdown('<div class="result-box">', unsafe_allow_html=True)
+                st.subheader("📝 Распознанный текст")
+                st.code(generated_text, language=None)
+                st.markdown('</div>', unsafe_allow_html=True)
+
+                st.download_button(
+                    label="💾 Скачать как .txt",
+                    data=generated_text,
+                    file_name="recognized_text.txt",
+                    mime="text/plain"
+                )
 
 st.markdown("---")
 st.caption("Сделано на базе [lightonai/LightOnOCR-1B-1025](https://huggingface.co/lightonai/LightOnOCR-1B-1025)")