app.py

import os
import re
import json
import base64
from io import BytesIO
from typing import Optional, List
from pathlib import Path

import torch
import numpy as np
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

# ─── Configuration ────────────────────────────────────────────────
MODEL_PATH = "fady-50/ocr-model"  # Your Hugging Face model ID
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ─── Load Model & Processor ──────────────────────────────────────
print(f"Loading model on {DEVICE}...")
processor = TrOCRProcessor.from_pretrained(MODEL_PATH, local_files_only=False)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_PATH, local_files_only=False)
model.to(DEVICE)
model.eval()
print("Model loaded successfully!")

# ─── FastAPI App ─────────────────────────────────────────────────
app = FastAPI(
    title="Prescription OCR API",
    description="Extract medicine names from handwritten prescriptions using TrOCR",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ─── Helper Functions ────────────────────────────────────────────
def preprocess_image(image: Image.Image) -> Image.Image:
    """Preprocess image for TrOCR."""
    # Convert to RGB if needed
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Resize to model input size (384x384 for TrOCR base)
    image = image.resize((384, 384), Image.Resampling.LANCZOS)
    return image


def extract_text_from_image(image: Image.Image) -> str:
    """Run OCR on image and return extracted text."""
    pixel_values = processor(images=image, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(DEVICE)

    # Generate text
    generated_ids = model.generate(
        pixel_values,
        max_length=512,
        num_beams=4,
        early_stopping=True,
        no_repeat_ngram_size=3,
    )

    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text.strip()


def extract_medicine_names(text: str) -> List[dict]:
    """
    Extract medicine names from OCR text.
    Uses pattern matching and common medical terms.
    """
    medicines = []

    # Common medicine name patterns (generic + brand names)
    # Pattern 1: Capitalized words that look like drug names (2-4 words)
    drug_pattern = r'(?:[A-Z][a-zA-Z]*(?:\s+[A-Z][a-zA-Z]*){0,3})'

    # Pattern 2: Words ending with common drug suffixes
    drug_suffixes = [
        'cin', 'mycin', 'micin', 'cillin', 'xacin', 'floxacin',
        'zole', 'sone', 'sone', 'nide', 'pril', 'sartan',
        'olol', 'dipine', 'zepam', 'zolam', 'pram', 'tidine',
        'tidine', 'mab', 'nib', 'vastatin', 'prastatin',
        'profen', 'coxib', 'triptan', 'platin', 'taxel'
    ]

    lines = text.split('\n')

    for line_num, line in enumerate(lines, 1):
        line = line.strip()
        if not line or len(line) < 3:
            continue

        # Check for drug suffixes
        has_drug_suffix = any(suffix.lower() in line.lower() for suffix in drug_suffixes)

        # Check if line looks like a medicine (starts with capital, contains dosage, or has drug suffix)
        dosage_indicators = ['mg', 'ml', 'g', 'mcg', 'IU', 'tablet', 'capsule', 
                            'syrup', 'suspension', 'injection', 'drops', 'cream',
                            'ointment', ' inhaler', 'spray', 'patch']

        has_dosage = any(ind.lower() in line.lower() for ind in dosage_indicators)

        # Score the line
        score = 0
        if has_drug_suffix:
            score += 3
        if has_dosage:
            score += 2
        if line[0].isupper() and len(line) > 3:
            score += 1
        if re.search(r'\d+', line):  # Contains numbers (likely dosage)
            score += 1

        # Extract potential medicine name (before dosage info)
        if score >= 2:
            # Try to extract just the medicine name (before numbers or dosage)
            name_match = re.match(r'^([A-Za-z][A-Za-z\s\-]+?)(?:\s+\d|\s*(?:mg|ml|g|\%|tablet|cap))', line, re.IGNORECASE)

            if name_match:
                med_name = name_match.group(1).strip()
            else:
                med_name = line

            # Clean up the name
            med_name = re.sub(r'[^A-Za-z\s\-]', '', med_name).strip()

            if len(med_name) > 2:
                medicines.append({
                    "name": med_name,
                    "full_line": line,
                    "line_number": line_num,
                    "confidence_score": min(score / 5.0, 1.0)
                })

    # Also try regex extraction for any missed medicines
    words = re.findall(r'[A-Z][a-zA-Z]{2,}(?:\s+[A-Z][a-zA-Z]{2,}){0,2}', text)
    existing_names = {m["name"].lower() for m in medicines}

    for word in words:
        word_clean = word.strip()
        if (word_clean.lower() not in existing_names and 
            len(word_clean) > 3 and
            any(suffix.lower() in word_clean.lower() for suffix in drug_suffixes)):
            medicines.append({
                "name": word_clean,
                "full_line": word_clean,
                "line_number": 0,
                "confidence_score": 0.6
            })

    # Sort by confidence
    medicines.sort(key=lambda x: x["confidence_score"], reverse=True)

    return medicines


# ─── API Endpoints ───────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "message": "Prescription OCR API",
        "status": "running",
        "model": "TrOCR Handwritten",
        "device": str(DEVICE)
    }


@app.get("/health")
def health_check():
    return {"status": "healthy", "model_loaded": True}


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    return_full_text: bool = Form(False),
    extract_medicines: bool = Form(True)
):
    """
    Upload a prescription image and extract text/medicine names.

    - **file**: Prescription image (JPG, PNG, etc.)
    - **return_full_text**: Return full OCR text
    - **extract_medicines**: Extract medicine names from text
    """
    # Validate file
    allowed_types = {"image/jpeg", "image/png", "image/jpg", "image/webp", "image/bmp"}
    if file.content_type not in allowed_types:
        raise HTTPException(
            status_code=400, 
            detail=f"Invalid file type. Allowed: {allowed_types}"
        )

    try:
        # Read and process image
        contents = await file.read()
        image = Image.open(BytesIO(contents))

        # Preprocess
        processed_image = preprocess_image(image)

        # Run OCR
        extracted_text = extract_text_from_image(processed_image)

        result = {
            "success": True,
            "filename": file.filename,
            "medicines_count": 0,
            "medicines": []
        }

        if return_full_text:
            result["full_text"] = extracted_text

        if extract_medicines:
            medicines = extract_medicine_names(extracted_text)
            result["medicines"] = medicines
            result["medicines_count"] = len(medicines)

        return JSONResponse(content=result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/predict_base64")
async def predict_base64(
    image_base64: str = Form(...),
    return_full_text: bool = Form(False),
    extract_medicines: bool = Form(True)
):
    """
    Send base64-encoded image and extract text/medicine names.
    """
    try:
        # Decode base64
        image_data = base64.b64decode(image_base64)
        image = Image.open(BytesIO(image_data))

        # Preprocess
        processed_image = preprocess_image(image)

        # Run OCR
        extracted_text = extract_text_from_image(processed_image)

        result = {
            "success": True,
            "medicines_count": 0,
            "medicines": []
        }

        if return_full_text:
            result["full_text"] = extracted_text

        if extract_medicines:
            medicines = extract_medicine_names(extracted_text)
            result["medicines"] = medicines
            result["medicines_count"] = len(medicines)

        return JSONResponse(content=result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/predict_url")
async def predict_url(
    image_url: str = Form(...),
    return_full_text: bool = Form(False),
    extract_medicines: bool = Form(True)
):
    """
    Provide image URL and extract text/medicine names.
    """
    import requests

    try:
        response = requests.get(image_url, timeout=30)
        response.raise_for_status()

        image = Image.open(BytesIO(response.content))
        processed_image = preprocess_image(image)
        extracted_text = extract_text_from_image(processed_image)

        result = {
            "success": True,
            "image_url": image_url,
            "medicines_count": 0,
            "medicines": []
        }

        if return_full_text:
            result["full_text"] = extracted_text

        if extract_medicines:
            medicines = extract_medicine_names(extracted_text)
            result["medicines"] = medicines
            result["medicines_count"] = len(medicines)

        return JSONResponse(content=result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ─── Gradio UI (for Hugging Face Spaces) ─────────────────────────
def create_gradio_interface():
    """Create Gradio interface for Hugging Face Spaces."""
    try:
        import gradio as gr

        def process_image(image):
            if image is None:
                return "Please upload an image", "[]"

            # Convert to PIL if needed
            if isinstance(image, np.ndarray):
                image = Image.fromarray(image)

            processed = preprocess_image(image)
            text = extract_text_from_image(processed)
            medicines = extract_medicine_names(text)

            # Format output
            medicines_text = "\n".join([
                f"{i+1}. {med['name']} (confidence: {med['confidence_score']:.0%})"
                for i, med in enumerate(medicines)
            ]) if medicines else "No medicines detected"

            return text, medicines_text

        demo = gr.Interface(
            fn=process_image,
            inputs=gr.Image(type="pil", label="Upload Prescription"),
            outputs=[
                gr.Textbox(label="Extracted Text", lines=10),
                gr.Textbox(label="Detected Medicines", lines=10)
            ],
            title="📋 Prescription OCR - Medicine Extractor",
            description="Upload a handwritten prescription image to extract medicine names.",
            examples=[]
        )
        return demo

    except ImportError:
        return None


# ─── Main Entry Point ────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn

    # Check if running in Hugging Face Spaces (Gradio)
    if os.environ.get("SPACE_ID") or os.environ.get("GRADIO_SERVER_NAME"):
        demo = create_gradio_interface()
        if demo:
            demo.launch(server_name="0.0.0.0", server_port=7860)
        else:
            uvicorn.run(app, host="0.0.0.0", port=7860)
    else:
        # Run FastAPI server
        port = int(os.environ.get("PORT", 7860))
        uvicorn.run(app, host="0.0.0.0", port=port)