Update app.py

app.py

@@ -6,8 +6,10 @@ from pydantic import BaseModel, ConfigDict
 import os
 from base64 import b64encode
 from io import BytesIO
-from PIL import Image  # Add this import
+from PIL import Image, ImageEnhance
 import logging
+import pytesseract
+import time
 
 app = FastAPI()
 
@@ -15,22 +17,18 @@ app = FastAPI()
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Get HuggingFace token from environment variable
 
 # Default model
 DEFAULT_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"
 
-class QueryRequest(BaseModel):
+class TextRequest(BaseModel):
     model_config = ConfigDict(protected_namespaces=())
-    
     query: str
-    image_data: Optional[str] = None  # Base64 encoded image data
     stream: bool = False
     model_name: Optional[str] = None
 
-class ChatForm(BaseModel):
+class ImageTextRequest(BaseModel):
     model_config = ConfigDict(protected_namespaces=())
-    
     query: str
     stream: bool = False
     model_name: Optional[str] = None
@@ -41,7 +39,7 @@ class ChatForm(BaseModel):
         query: str = Form(...),
         stream: bool = Form(False),
         model_name: Optional[str] = Form(None),
-        image: Optional[UploadFile] = File(None)
+        image: UploadFile = File(...)  # Make image required for i2t2t
     ):
         return cls(
             query=query,
@@ -52,9 +50,7 @@ class ChatForm(BaseModel):
 def get_client(model_name: Optional[str] = None):
     """Get inference client for specified model or default model"""
     try:
-        # Use provided model_name if it exists and is not empty, otherwise use DEFAULT_MODEL
         model_path = model_name if model_name and model_name.strip() else DEFAULT_MODEL
-        
         return InferenceClient(
             model=model_path
         )
@@ -64,26 +60,11 @@ def get_client(model_name: Optional[str] = None):
             detail=f"Error initializing model {model_path}: {str(e)}"
         )
 
-def generate_response(query: str, image_data: Optional[str] = None, model_name: Optional[str] = None):
-    messages = []
-    
-    # Create the system and user message
-    user_content = f"[SYSTEM] You are ASSISTANT who answer question asked by user in short and concise manner. [USER] {query}"
-    
-    # If there's an image, add it to the message
-    if image_data:
-        messages.append({
-            "role": "user",
-            "content": [
-                {"type": "text", "text": user_content},
-                {"type": "image_url", "image_url": {"url": f"data:image/*;base64,{image_data}"}}
-            ]
-        })
-    else:
-        messages.append({
-            "role": "user",
-            "content": user_content
-        })
+def generate_text_response(query: str, model_name: Optional[str] = None):
+    messages = [{
+        "role": "user",
+        "content": f"[SYSTEM] You are ASSISTANT who answer question asked by user in short and concise manner. [USER] {query}"
+    }]
 
     try:
         client = get_client(model_name)
@@ -97,61 +78,137 @@ def generate_response(query: str, image_data: Optional[str] = None, model_name:
     except Exception as e:
         yield f"Error generating response: {str(e)}"
 
+def generate_image_text_response(query: str, image_data: str, model_name: Optional[str] = None):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"[SYSTEM] You are ASSISTANT who answer question asked by user in short and concise manner. [USER] {query}"},
+                {"type": "image_url", "image_url": {"url": f"data:image/*;base64,{image_data}"}}
+            ]
+        }
+    ]
+    
+    logger.debug(f"Messages sent to API: {messages}")
+    
+    try:
+        client = get_client(model_name)
+        for message in client.chat_completion(messages, max_tokens=2048, stream=True):
+            logger.debug(f"Received message chunk: {message}")
+            token = message.choices[0].delta.content
+            yield token
+    except Exception as e:
+        logger.error(f"Error in generate_image_text_response: {str(e)}")
+        yield f"Error generating response: {str(e)}"
+
+def preprocess_image(img):
+    """Enhance image for better OCR results"""
+    # Convert to grayscale
+    img = img.convert('L')
+    
+    # Enhance contrast
+    enhancer = ImageEnhance.Contrast(img)
+    img = enhancer.enhance(2.0)
+    
+    # Enhance sharpness
+    enhancer = ImageEnhance.Sharpness(img)
+    img = enhancer.enhance(1.5)
+    
+    return img
+
 @app.get("/")
 async def root():
     return {"message": "Welcome to FastAPI server!"}
 
-@app.post("/chat")
-async def chat(form_data: tuple[ChatForm, Optional[UploadFile]] = Depends(ChatForm.as_form)):
+@app.post("/t2t")
+async def text_to_text(request: TextRequest):
+    try:
+        if request.stream:
+            return StreamingResponse(
+                generate_text_response(request.query, request.model_name),
+                media_type="text/event-stream"
+            )
+        else:
+            response = ""
+            for chunk in generate_text_response(request.query, request.model_name):
+                response += chunk
+            return {"response": response}
+    except Exception as e:
+        logger.error(f"Error in /t2t endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/i2t2t")
+async def image_text_to_text(form_data: tuple[ImageTextRequest, UploadFile] = Depends(ImageTextRequest.as_form)):
     form, image = form_data
     try:
-        image_data = None
-        if image:
-            logger.debug("Image received")
-            # Read the image
-            contents = await image.read()
+        # Process image
+        contents = await image.read()
+        try:
+            logger.debug("Attempting to open image")
+            img = Image.open(BytesIO(contents))
+            if img.mode != 'RGB':
+                img = img.convert('RGB')
             
-            # Convert image to appropriate format if needed
-            try:
-                logger.debug("Attempting to open image")
-                img = Image.open(BytesIO(contents))
-                logger.debug(f"Image format before conversion: {img.format}, mode: {img.mode}")
-                # Convert to RGB if needed
-                if img.mode != 'RGB':
-                    img = img.convert('RGB')
-                logger.debug(f"Image format after conversion: {img.format}, mode: {img.mode}")
-                
-                # Save as JPEG in memory
-                buffer = BytesIO()
-                img.save(buffer, format="JPEG")
-                image_data = b64encode(buffer.getvalue()).decode('utf-8')
-                logger.debug("Image processed and encoded to base64")
-            except Exception as img_error:
-                logger.error(f"Error processing image: {str(img_error)}")
-                raise HTTPException(
-                    status_code=422,
-                    detail=f"Error processing image: {str(img_error)}"
-                )
+            buffer = BytesIO()
+            img.save(buffer, format="PNG")
+            image_data = b64encode(buffer.getvalue()).decode('utf-8')
+            logger.debug("Image processed and encoded to base64")
+        except Exception as img_error:
+            logger.error(f"Error processing image: {str(img_error)}")
+            raise HTTPException(
+                status_code=422,
+                detail=f"Error processing image: {str(img_error)}"
+            )
         
         if form.stream:
             return StreamingResponse(
-                generate_response(form.query, image_data, form.model_name),
+                generate_image_text_response(form.query, image_data, form.model_name),
                 media_type="text/event-stream"
             )
         else:
             response = ""
-            for chunk in generate_response(form.query, image_data, form.model_name):
+            for chunk in generate_image_text_response(form.query, image_data, form.model_name):
                 response += chunk
             return {"response": response}
     except Exception as e:
-        logger.error(f"Error in /chat endpoint: {str(e)}")
+        logger.error(f"Error in /i2t2t endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(
-        "main:app", 
-        port=8000,
-        reload=True,  # Enable auto-reload
-        reload_dirs=["./"]  # Watch the current directory for changes
-    )
+@app.post("/tes")
+async def ocr_endpoint(image: UploadFile = File(...)):
+    try:
+        # Read and process the image
+        contents = await image.read()
+        img = Image.open(BytesIO(contents))
+        
+        # Preprocess the image
+        img = preprocess_image(img)
+        
+        # Perform OCR with timeout and retries
+        max_retries = 3
+        text = ""
+        
+        for attempt in range(max_retries):
+            try:
+                text = pytesseract.image_to_string(
+                    img,
+                    timeout=30,  # 30 second timeout
+                    config='--oem 3 --psm 6'
+                )
+                break
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    raise HTTPException(
+                        status_code=500,
+                        detail=f"Error extracting text: {str(e)}"
+                    )
+                time.sleep(1)  # Wait before retry
+        
+        return {"text": text}
+        
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error processing image: {str(e)}"
+        )
+