ococtata/Software-Engineering-Project-OCR-AI

1from fastapi import FastAPI, UploadFile, File, HTTPException2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import JSONResponse4from google import genai5from PIL import Image6import io7import json8import re9import os10from dotenv import load_dotenv11 12load_dotenv()13 14app = FastAPI()15 16app.add_middleware(17    CORSMiddleware,18    allow_origins=["*"], 19    allow_credentials=True,20    allow_methods=["*"],21    allow_headers=["*"],22)23 24api_key = os.getenv("GENAI_API_KEY")25if not api_key:26    raise ValueError("GENAI_API_KEY not found in environment variables")27 28client = genai.Client(api_key=api_key)29 30prompt = (31    "Transcribe the nutrition facts table in this image and output the data "32    "as a single **JSON object**. Use keys like 'serving-size', 'energy-kcal', 'fat', 'carbohydrates', 'proteins', 'saturated-fat', 'trans-fat', 'sugars', 'added-sugars', 'sodium', 'salt', and 'fiber'. "33    "Do not include any text outside of the JSON object."34    "If the values are not in the image, fill the values with 0"35    "Pay attention to the unit, normalize the unit so it's stated in g (gram) instead of mg (miligram)"36)37 38@app.post("/ocr")39async def ocr(file: UploadFile = File(...)):40    try:41        print(f"Received file: {file.filename}")42        43        img_bytes = await file.read()44        image = Image.open(io.BytesIO(img_bytes))45        print(f"Image loaded: {image.size}, format: {image.format}")46 47        print("Calling Gemini API...")48        result = client.models.generate_content(49            model='gemini-2.0-flash-exp', 50            contents=[prompt, image]51        )52 53        data = result.text54        print(f"Raw Gemini response: {data}")55 56        cleanedData = data.strip().replace('```json', '').replace('```', '').replace('\n','').strip()57        print(f"Cleaned step 1: {cleanedData}")58 59        cleanedData = re.sub(r'\s+', ' ', cleanedData).strip()60        print(f"Cleaned step 2: {cleanedData}")61 62        tempDict = {}63 64        for i in cleanedData.split(','):65            if ':' not in i:66                continue67            x = i.split(':')68            try:69                col = (x[0].split('"'))[1]70            except:71                col = x[0].replace('{', '').replace('"', '').strip()72            num = x[1]73 74            tempDict[col] = num75            print(f"Parsed: {col} = {num}")76 77        print(f"tempDict: {tempDict}")78 79        divider = None80        for i in tempDict:81            if i == 'serving-size':82                divider_raw = tempDict[i].split('"')83                divider_str = divider_raw[1].split('g')84                divider = int(divider_str[0].strip())85                print(f"Found serving-size divider: {divider}g")86                break87 88        if divider is None or divider == 0:89            print("Warning: No valid serving-size found, defaulting to 1")90            divider = 191 92        resDict = {}93 94        for i in tempDict:95            if i == 'serving-size':96                continue97            else:98                try:99                    value = float(tempDict[i].strip())100                    resDict[i+'_1g'] = round(value/divider, 3)101                except ValueError:102                    try:103                        value_str = tempDict[i].split("}")104                        value = float(value_str[0].strip())105                        resDict[i+'_1g'] = round(value/divider, 3)106                    except:107                        print(f"Warning: Could not parse value for {i}, setting to 0")108                        resDict[i+'_1g'] = 0109 110        print(f"Final result: {resDict}")111        return JSONResponse(resDict)112 113    except Exception as e:114        print(f"Error in OCR: {str(e)}")115        import traceback116        traceback.print_exc()117        raise HTTPException(status_code=500, detail=str(e))118 119@app.get("/")120def home():121    return {122        "status": "running",123        "endpoints": {124            "/": "Health check",125            "/ocr": "POST - Extract nutrition data from image (returns per 1g values)"126        }127    }128 129@app.get("/debug-key")130def debug_key():131    return {"key": os.getenv("GENAI_API_KEY")}132 133if __name__ == "__main__":134    import uvicorn135    uvicorn.run(app, host="0.0.0.0", port=7860)