stephenz007/pdf-api

1import requests2import json3import pdfplumber4import pandas as pd5import time6from cnocr import CnOcr7from sentence_transformers import SentenceTransformer, models, util8# import mysql.connector9 10from fastapi import FastAPI, UploadFile, File11 12from fastapi.responses import HTMLResponse13 14from fastapi.middleware.cors import CORSMiddleware15 16app = FastAPI()17 18 19origins = [20    "http://localhost",21    "http://localhost:8080",22    "http://localhost:3000",23    "https://invoice-pdf-xi.vercel.app"24]25 26app.add_middleware(27    CORSMiddleware,28    allow_origins=origins,29    allow_credentials=True,30    allow_methods=["*"],31    allow_headers=["*"],32)33 34@app.get("/")35def home():36    html_content = open('index.html').read()37    return HTMLResponse(content=html_content, status_code=200)38 39 40@app.post("/file")41async def up_file(file: UploadFile = File(...)):42    ocr = CnOcr()43    doc_text_list = []44    with pdfplumber.open(file.file) as pdf:45        for i in range(len(pdf.pages)):46            # Read page i+1 of a PDF document47            page = pdf.pages[i]48            res_list = page.extract_text().split('\n')[:-1]49 50            for j in range(len(page.images)):51                # Get the binary stream of the image52                img = page.images[j]53                file_name = '{}-{}-{}.png'.format(str(time.time()), str(i), str(j))54                with open(file_name, mode='wb') as f:55                    f.write(img['stream'].get_data())56                try:57                    res = ocr.ocr(file_name)58                except Exception as e:59                    res = []60                if len(res) > 0:61                    res_list.append(' '.join([re['text'] for re in res]))62 63            tables = page.extract_tables()64            for table in tables:65                # The first column is used as the header66                df = pd.DataFrame(table[1:], columns=table[0])67                try:68                    records = json.loads(df.to_json(orient="records", force_ascii=False))69                    for rec in records:70                        res_list.append(json.dumps(rec, ensure_ascii=False))71                except Exception as e:72                    res_list.append(str(df))73 74            doc_text_list += res_list75    # doc_text_list = [str(text).strip() for text in doc_text_list if len(str(text).strip()) > 0]76    print(doc_text_list)77    return doc_text_list78 79# @app.get("/{provider_id}")80# def get_provider(provider_id):81    conn = mysql.connector.connect(82        host="localhost",83        user="root",84        passwd="1234",85        database="cg_app",86        auth_plugin='mysql_native_password'87    )88 89    if conn.is_connected():90        print('Connected to MySQL database')91 92    cursor = conn.cursor()93 94    query = f"SELECT * FROM email_automation_invoice_provider_profile_updated WHERE provider_id={provider_id}"95    cursor.execute(query)96 97    # Get column names from cursor description98    column_names = [col[0] for col in cursor.description]99 100    # Create list of dictionaries representing rows101    rows = []102    for row in cursor.fetchall():103        row_dict = {}104        for i, value in enumerate(row):105            row_dict[column_names[i]] = value106        rows.append(row_dict)107 108    cursor.close()109    conn.close()110 111    # Return result as JSON response112    return rows