Upload 3 files

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \ 
+    libgl1-mesa-glx \ 
+    libglib2.0-0 \ 
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy model files and app
+COPY . .
+
+# Expose port
+EXPOSE 7860
+
+# Run the app
+CMD ["python", "main.py"]

main.py

@@ -0,0 +1,372 @@
+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)

requirements.txt

@@ -0,0 +1,10 @@
+transformers>=4.30.0
+torch>=2.0.0
+Pillow>=9.5.0
+fastapi>=0.100.0
+uvicorn[standard]>=0.23.0
+python-multipart>=0.0.6
+accelerate>=0.20.0
+numpy>=1.24.0
+gradio>=4.0.0
+requests