main.py

import cv2
import pytesseract
import os
import re
from rapidfuzz import fuzz
from fastapi import FastAPI, UploadFile, Form
from fastapi.responses import JSONResponse
import uvicorn
from PIL import Image, ExifTags,ImageOps
import numpy as np

app = FastAPI()
output_dir = "hasil"
os.makedirs(output_dir, exist_ok=True)
gambar_dir = "gambar"
os.makedirs(gambar_dir, exist_ok=True)


def load_image_fix_orientation(path: str) -> str:
    try:
        img = Image.open(path)
        img = ImageOps.exif_transpose(img)
        img.save(path)
        return path
    except Exception as e:
        print("Orientation fix error:", e)
        return path

def ocr_raw(image, suffix=""):
    custom_config = r'--oem 3 --psm 6'
    text = pytesseract.image_to_string(image, config=custom_config)
    cv2.imwrite(f"{output_dir}/step0_raw{suffix}.jpg", image)
    return text

def ocr_preprocessed(image, suffix=""):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    enhanced = cv2.convertScaleAbs(gray, alpha=1.2, beta=0)
    cv2.imwrite(f"{output_dir}/step0_raw.jpg", enhanced)
    custom_config = r'--oem 3 --psm 6'
    text = pytesseract.image_to_string(enhanced, config=custom_config)
    return text

def ocr_preprocessed2(image, suffix=""):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    enhanced = cv2.convertScaleAbs(gray, alpha=1.5, beta=-20)
    cv2.imwrite(f"{output_dir}/step0_raw.jpg", enhanced)
    custom_config = r'--oem 3 --psm 6'
    text = pytesseract.image_to_string(enhanced, config=custom_config)
    return text

def normalize_candidates(candidates: list[str]) -> list[str]:
    result = []
    for c in candidates:
        match = re.match(r"([A-Za-z]+)(\d.*)", c)
        if match:
            letters, rest = match.groups()
            # sisain max 2 huruf aja
            fixed = letters[:1] + rest
            result.append(fixed)
        else:
            # kalau ga ada huruf depannya (langsung angka) → tetap
            result.append(c)
    return result
  
def check_similarity(text, keyword, threshold=85):
    filtered = re.sub(r"[^a-zA-Z0-9\n]", "", text)

    candidates = filtered.split("\n")
    candidates =normalize_candidates(candidates)
    cleaned_candidates = [re.sub(r"\s+", "", line) for line in candidates if line.strip()]
    #print(cleaned_candidates)
    for cand in cleaned_candidates:
        score = fuzz.ratio(cand, keyword)
        if score >= threshold:
            return True, cand, score
    return False, None, 0

# --- Step 0: pastikan horizontal ---


def process_image(image_path: str, keyword: str):
    img = cv2.imread(image_path)
    h, w = img.shape[:2]
    if h > w:  # portrait
        print("📐 Gambar portrait → rotate ke landscape (90°)")
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)

    # --- Step 1: OCR langsung ---
    found_text = ocr_raw(img)
    ok, match, score = check_similarity(found_text, keyword)

    # --- Step 2: Preprocessing ---
    if not ok:
        found_text = ocr_preprocessed2(img)
        ok, match, score = check_similarity(found_text, keyword)

    if not ok:
        found_text = ocr_preprocessed(img)
        ok, match, score = check_similarity(found_text, keyword)

    # --- Step 3: Rotate 180° (tetap landscape) ---
    if not ok:
        img_rot = cv2.rotate(img, cv2.ROTATE_180)
        print("📐 Gambar portrait → rotate ke landscape (180°)")
        found_text = ocr_preprocessed(img_rot, suffix="_rot180")
        ok, match, score = check_similarity(found_text, keyword)
        img = img_rot

    # --- Simpan hasil akhir ---
    cv2.imwrite(f"{output_dir}/step0_raw.jpg", img)
    with open(f"{output_dir}/hasil_ocr.txt", "w", encoding="utf-8") as f:
        f.write(found_text)

    print("match",match)
    print("keyword",keyword)
    return {
        "found": ok,
        "match": keyword,#match,
        "score": score
    }


# --- API Endpoint ---
@app.post("/ocr")
async def ocr_endpoint(file: UploadFile, keyword: str = Form(...)):
    file_path = os.path.join(gambar_dir, file.filename)
    contents = await file.read()
    with open(file_path, "wb") as f:
        f.write(contents)

    image = load_image_fix_orientation(file_path)
    
    result = process_image(image, keyword)
    return JSONResponse(content=result)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8124)