Upload main.py

main.py

@@ -0,0 +1,192 @@
+from unittest import result
+
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from google import genai
+from httpcore import request
+from google.genai import types
+#from dotenv import load_dotenv
+from pydantic import BaseModel
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import cv2
+from paddleocr import PaddleOCR
+import numpy as np
+import time
+import os
+
+os.environ["OMP_NUM_THREADS"] = "1"
+
+MODEL_ID = "models/gemini-2.5-flash"
+MODEL_FALLBACK = "models/gemini-3.1-flash-lite-preview"
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+modelo_local = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
+#load_dotenv()
+
+api_key = os.environ.get("GEMINI_API_KEY")
+
+client = genai.Client(api_key=api_key)
+
+
+config = types.GenerateContentConfig(
+    temperature=0,
+)
+
+class ChatRequest(BaseModel):
+    pregunta: str
+    contexto: str = ""
+    historial: list[str] = []
+
+class ChatResponse(BaseModel):
+    respuesta: str
+    historial: list[str] = []
+
+class OCRResponse(BaseModel):
+	texto: str
+
+app = FastAPI(title="OCR API", version="1.0.0")
+ocr = PaddleOCR(use_angle_cls=True, lang="es")
+
+try:
+    local_model = pipeline("text2text-generation", model="google/flan-t5-small")
+except Exception as e:
+    print(f"Aviso: No se pudo cargar el modelo local: {e}")
+    local_model = None
+
+app.mount("/static", StaticFiles(directory="static", html=True), name="static")
+
+@app.get("/")
+def home():
+   return FileResponse("static/index.html")
+
+
+
+@app.post("/api/chat", response_model=ChatResponse)
+def prompt(request: ChatRequest):
+    max_reintentos = 2
+    segundos_espera = 1.5
+    e = "Error desconocido"
+
+    try:
+        pregunta = request.pregunta
+        historial = request.historial
+        texto_ocr = request.contexto  # <--- Tomamos el texto enviado por el frontend
+
+        texto_prompt = f"""
+                Eres un asistente que responde preguntas usando SOLO información del documento.
+
+                REGLAS IMPORTANTES:
+                - No copies el documento completo.
+                - No repitas texto largo del documento.
+                - Extrae SOLO la información necesaria.
+                - Si la respuesta no está en el documento, di: "No aparece en el documento".
+                - Responde de forma breve y directa.
+
+                DOCUMENTO:
+                \"\"\"{texto_ocr}\"\"\"  <--- Cambiado: ahora usa el OCR
+
+                HISTORIAL:
+                {historial}
+
+                PREGUNTA:
+                {pregunta}
+
+                RESPUESTA:
+                """
+
+        for intento in range(max_reintentos):
+            try:
+                response = client.models.generate_content(
+                        model=MODEL_ID,
+                        contents=texto_prompt,
+                        config=config
+                )
+                return ChatResponse(respuesta=response.text)
+            except Exception as ex:
+                e = ex
+                print(f"Intento {intento+1} fallido: {e}")
+                time.sleep(segundos_espera)
+                
+        try:
+            response = client.models.generate_content(
+                model=MODEL_FALLBACK,
+                contents=texto_prompt,
+                config=config
+            )
+            return ChatResponse(respuesta=response.text)
+        except Exception as e_fallback:
+            print(f"Fallido: {e_fallback}")
+            try:
+            # Simplificamos el prompt para el modelo local pequeño
+                res_local = modelo_local(f"question: {pregunta} context: {texto_ocr}", max_new_tokens=50)
+                return ChatResponse(respuesta=res_local[0]['generated_text'])
+            except Exception as e_final:
+                raise HTTPException(status_code=500, detail="Error en todos los modelos (incluyendo local)")
+
+    except Exception as e:
+        if isinstance(e, HTTPException): raise e
+        raise HTTPException(status_code=500, detail=str(e))
+
+    except Exception as e:
+        if isinstance(e, HTTPException): raise e
+        raise HTTPException(status_code=500, detail=str(e))
+    
+
+def extraer_texto(resultado_ocr):
+    textos = []
+    
+    # PaddleOCR devuelve una lista de páginas. 
+    # Validamos que el resultado no sea None y que la primera página tenga contenido.
+    if resultado_ocr and resultado_ocr[0] is not None:
+        for linea in resultado_ocr[0]:
+            # linea[1][0] es donde reside el texto detectado
+            textos.append(linea[1][0])
+
+    return " ".join(textos)
+
+      
+def preprocesar_imagen(image_bytes):
+    # Convertir bytes a imagen
+    nparr = np.frombuffer(image_bytes, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+    if img is None:
+        raise ValueError("No se pudo decodificar la imagen")
+
+    # PaddleOCR maneja internamente el binarizado, 
+    # es mejor enviarle la imagen limpia o solo en gris.
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    return gray
+    
+@app.post("/api/ocr", response_model=OCRResponse)
+async def ocr_image(file: UploadFile = File(...)):
+    if file.content_type not in {"image/png", "image/jpeg", "image/jpg", "image/webp"}:
+        raise HTTPException(status_code=400, detail="Formato no soportado")
+
+    try:
+        # Leemos los bytes del archivo
+        contenido = await file.read()
+        
+        # Preprocesamos y ejecutamos OCR
+        img = preprocesar_imagen(contenido)
+        
+        # cls=True activa la clasificación de ángulo si se configuró en la instancia
+        result = ocr.ocr(img, cls=True) 
+        
+        texto_extraido = extraer_texto(result)
+        
+        if not texto_extraido.strip():
+            return OCRResponse(texto="No se detectó texto en la imagen.")
+            
+        return OCRResponse(texto=texto_extraido)
+
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error OCR: {str(e)}")
+
+    
+    
+