from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse from fastapi import File, UploadFile from fastapi.responses import StreamingResponse from typing import List from pdf2image import convert_from_bytes import csv import io from transformers import pipeline app = FastAPI() """## Poppler dir""" poppler_path = "poppler-23.11.0/Library/bin" @app.post("/classify") async def classify_doc(files: List[UploadFile] = File(...)): classificationResults = {} for file in files: try: contents = file.file.read() filename = file.filename if filename.endswith('.pdf'): try: pages = convert_from_bytes(contents, poppler_path = poppler_path) print(pages) for pagenum, image in enumerate(pages): classificationRes, dtype_conf = doctype_classify(image.convert('RGB'), filename) # add/update classification result dictionary if (classificationRes in classificationResults): classificationResults.update({classificationRes : classificationResults[classificationRes] + 1}) else: classificationResults.update({classificationRes : 1}) except Exception as err: print(err) return f"Error in opening {filename}, {err}" # png, jpg, jpeg files else: classificationRes = classify_acct_dtype_str(contents, filename) # add/update classification result dictionary if (classificationRes in classificationResults): classificationResults.update({classificationRes : classificationResults[classificationRes] + 1}) else: classificationResults.update({classificationRes : 1}) except Exception as err: print(Exception, err) return {"message": "There was an error in uploading file(s)"} finally: file.file.close() # Convert dictionary to CSV string csv_data = io.StringIO() csv_writer = csv.writer(csv_data) csv_writer.writerow(["Type", "Count"]) # Header row for key, value in classificationResults.items(): csv_writer.writerow([key, value]) return StreamingResponse( iter([csv_data.getvalue()]), media_type="text/csv", headers={"Content-Disposition": f"attachment; filename=data.csv"} ) # return {"message": f"{[file.filename for file in files]} : {[classifyFiles(file) for file in files]}"} def classifyFiles(file): try: contents = file.file.read() filename = file.filename classificationResults = [] if filename.endswith('.pdf'): try: pages = convert_from_bytes(open(file, 'rb').read()) for pagenum, image in enumerate(pages): if pagenum != 0 and pagenum < len(pages): classificationRes = classify_acct_dtype_str(contents, filename) # classificationResults[f"{pagenum:02d}"] = { # 'doctype': classificationRes # } except: return f"Error in opening {filename}" else: classificationRes = classify_acct_dtype_str(contents, filename) # classificationResults[f"{0:02d}"] = { # 'doctype' : classificationRes # } except Exception as err: print(Exception, err) return {"message": "There was an error in uploading file(s)"} finally: file.file.close() return classificationRes app.mount("/", StaticFiles(directory="static", html=True), name="static") @app.get("/") def index() -> FileResponse: return FileResponse(path="/app/static/index.html", media_type="text/html") import re import torch from transformers import DonutProcessor, VisionEncoderDecoderModel from datasets import load_dataset import os from PIL import Image # Doc classifier model classifier_doctype_processor = DonutProcessor.from_pretrained("calumpianojericho/donutclassifier_acctdocs_by_doctype") classifier_doctype_model = VisionEncoderDecoderModel.from_pretrained("calumpianojericho/donutclassifier_acctdocs_by_doctype") """### Inference Code""" def inference(input, model, processor, threshold=1.0, task_prompt="", get_confidence=False): device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) is_confident = True decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids pil_img=input image = np.array(pil_img) pixel_values = processor(image, return_tensors="pt").pixel_values outputs = model.generate( pixel_values.to(device), decoder_input_ids=decoder_input_ids.to(device), max_length=model.decoder.config.max_position_embeddings, early_stopping=True, pad_token_id=processor.tokenizer.pad_token_id, eos_token_id= processor.tokenizer.eos_token_id, use_cache=True, num_beams=1, bad_words_ids=[[processor.tokenizer.unk_token_id]], return_dict_in_generate=True, output_scores=True, ) sequence = processor.batch_decode(outputs.sequences)[0] sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "") sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() # remove first task start token seq = processor.token2json(sequence) if get_confidence: return seq, pred_confidence(outputs.scores, threshold) return seq def pred_confidence(output_scores, threshold): is_confident=True for score in output_scores: exp_scores = np.exp(score[0].cpu().numpy()) # scores are logits, we use the exp function so that all values are positive sum_exp = np.sum(exp_scores) # taking the sum of the token scores idx = np.argmax(exp_scores) # taking the index of the token with the highest score prob_max = exp_scores[idx]/sum_exp # normalizing the token with the highest score wrt the sum of all scores. Returns probability if prob_max < threshold: is_confident = False # print(prob_max) return is_confident CUDA_LAUNCH_BLOCKING=1 def parse_text(input, filename): model = base_model processor = base_processor seq = inference(input, model, processor, task_prompt="") return str(seq) def doctype_classify(input, filename): model = classifier_doctype_model processor = classifier_doctype_processor seq, is_confident = inference(input, model, processor, threshold=0.90, task_prompt="", get_confidence=True) return seq.get('class'), is_confident def account_classify(input, filename): model = classifier_account_model processor = classifier_account_processor seq, is_confident = inference(input, model, processor, threshold=0.999, task_prompt="", get_confidence=True) return seq.get('class'), is_confident """## Text processing/string matcher code""" import locale locale.getpreferredencoding = lambda: "UTF-8" """## Classify Document Images""" import numpy as np import csv import re import os import requests from io import BytesIO def classify_acct_dtype_str(content, filename): # response = requests.get("https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/city-streets.jpg") # ipt = Image.open(BytesIO(response.content)) try: ipt = Image.open(BytesIO(content)) dtype_inf, dtype_conf = doctype_classify(ipt, filename) except Exception as err: return f"Error in opening {filename}, {err}" return dtype_inf # classify_acct_dtype_str("https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/city-streets.jpg")