our-sci/agriculture-experiments-document-processing
1
1import gradio as gr2import pandas as pd3import re4 5from PIL import Image, ImageDraw, ImageFont6import torch7from transformers import LayoutLMv3Processor, LayoutLMv3ForQuestionAnswering, LayoutLMv3ForTokenClassification8 9processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")10 11# More traditional approach that works from token classification basis (not questions)12model = LayoutLMv3ForTokenClassification.from_pretrained("microsoft/layoutlmv3-base")13device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')14 15print_device_name = torch.cuda.get_device_name(torch.cuda.current_device())16 17print(f"Debug -- Using device: {device}")18print(f"Debug -- Current Device Name: {print_device_name}")19model.to(device)20 21labels = model.config.id2label22print(labels)23 24# Homemade feature extraction 25def extract_features(tokens, labels): 26 merged_entities = []27 current_date = ""28 29 print(f"Debug -- Starting entity extraction")30 #date_pattern = r"\d{1,2}/\d{1,2}/\d{2,4}" # Matches full date formats like MM/DD/YYYY or DD/MM/YYYY31 #partial_date_pattern = r"\d{1,2}$|[/-]$" # Matches partial date components like "12" or "/" at the end32 33 #date_pattern = r"\d{1,2}/\d{1,2}/\d{2,4}" # Matches full date formats like MM/DD/YYYY or DD/MM/YYYY34 #partial_date_pattern = r"^\d{1,2}/?$|^[/-]$" # Matches partial date components like "12", "/", "02/", etc.35 36 date_pattern = r"^\d{2}/\d{2}/\d{2}(\d{2})?$"37 partial_date_pattern = r"^\d{1,2}/?$|^/$"38 39 # This is a label AGNOSTIC approach 40 for token, label in zip(tokens, labels): 41 print(f"Debug -- Processing token: {token}")42 43 # If we already have some part of a date and the next token could still be part of it, continue accumulating44 if current_date and re.match(partial_date_pattern, token): 45 current_date += token46 print(f"Debug -- Potential partial date: {current_date}")47 # If the accumulated entity matches a complete date after appending this token48 elif re.match(date_pattern, current_date + token):49 current_date += token50 merged_entities.append((current_date, 'date'))51 print(f"Debug -- Complete date added: {current_date}")52 current_date = "" # Reset for next entity53 # If the token could start a new date (e.g., '14' could be a day or hour)54 elif re.match(partial_date_pattern, token):55 current_date = token56 print(f"Debug -- Potentially starting a new date: {token}")57 else: 58 # If no patterns are detected and there is any accumulated data59 #if current_date: 60 # # Finalize accumulated partial date61 # print(f"Debug -- Date finalized: {current_date}")62 # merged_entities.append((current_date, 'date'))63 # current_date = "" # Reset for next entity64 65 # Append token as non-date66 print(f"Debug -- Appending non-date Token: {token}")67 merged_entities.append((token, 'non-date'))68 69 # If there's any leftover accumulated date data, add it to merged_entities70 if current_date:71 print(f"Debug -- Dangling leftover date added: {current_date}")72 merged_entities.append((current_date, 'date'))73 74 return merged_entities75 76 77 78 79 # NOTE: labels aren't being applied properly ... This is the LABEL approach 80 #81 # Loop through tokens and labels 82 #for token, label in zip(tokens, labels): 83 # print(f"Debug -- Potentially creating date,, token: {token} label: {label}")84 # 85 # if label == 'LABEL_1':86 # # Check for partial date fragments (like '12' or '/')87 # if re.match(date_pattern, current_date):88 # merged_entities.append((current_date, 'date'))89 # print(f"Debug -- Complete date added: {token}")90 # current_date = "" # Reset for next entity91 # # If the accumulated entity matches a full date92 # elif re.match(partial_date_pattern, token):93 # print(f"Debug -- Potentially building date: Token Start {token} After Token")94 # current_date += token # Append token to the current entity95 # else: 96 # # No partial or completed patterns are detected, but it's still LABEL_197 # # If there were any accumulated data so far98 # if current_date: 99 # merged_entities.append((current_date, 'date'))100 # print(f"Debug -- Date finalized: {current_date}")101 # current_date = "" # Reset102 # 103 # merged_entities.append((token, label))104 # else: 105 # # These are LABEL_0, supposedly trash but keep them for now106 # if current_date: # If there's a leftover date fragment, add it first107 # merged_entities.append((current_date, 'date'))108 # print(f"Debug -- Finalizing leftover date added: {current_date}")109 # current_date = "" # Reset110#111# # Append LABEL_0 token112# print(f"Debug -- Appending LABEL_0 Token: Token Start {token} Token After")113# merged_entities.append((token, label))114#115# if current_date:116# print(f"Debug -- Dangling leftover date added: {current_date}")117# merged_entities.append((current_date, 'date'))118#119# return merged_entities120 121 122# process the image in the correct format123# extract token classifications 124def parse_ticket_image(image): 125 126 # Process image127 if image: 128 document = image.convert("RGB") if image.mode != "RGB" else image129 else: 130 print(f"Warning - no image or malformed image!") 131 return pd.DataFrame()132 133 # Encode document image134 encoding = processor(document, return_tensors="pt", truncation=True)135 136 # Move encoding to appropriate device137 for k, v in encoding.items(): 138 encoding[k] = v.to(device)139 140 # Perform inference141 outputs = model(**encoding)142 143 # extract predictions144 predictions = outputs.logits.argmax(-1).squeeze().tolist()145 146 input_ids = encoding.input_ids.squeeze().tolist()147 words = [processor.tokenizer.decode(id) for id in input_ids]148 149 extracted_fields = []150 151 for idx, pred in enumerate(predictions): 152 label = model.config.id2label[pred]153 extracted_fields.append((label, words[idx]))154 # apparently stands for non-entity tokens155 #if label != 'LABEL_0' and '<' not in words[idx]: 156 # extracted_fields.append((label, words[idx]))157 158 if len(extracted_fields) == 0:159 print(f"Warning - no fields were extracted!") 160 return pd.DataFrame(columns=["Field", "Value"])161 162 # Create lists for fields and values 163 fields = [field[0] for field in extracted_fields]164 values = [field[1] for field in extracted_fields]165 166 # Ensure both lists have the same length167 min_length = min(len(fields), len(values))168 fields = fields[:min_length]169 values = values[:min_length]170 171 #Homemade feature extraction 172 values = extract_features(values, fields)173 174 #Ensure both lists have the same length175 min_length = min(len(fields), len(values))176 fields = fields[:min_length]177 values = values[:min_length]178 179 data = {180 "Field": fields,181 "Value": values182 }183 df = pd.DataFrame(data)184 185 return df186 187 188# This is how to use questions to find answers in the document189# Less traditional approach, less flexibility, easier to implement/understand (didnt provide robust answers)190#model = LayoutLMv3ForQuestionAnswering.from_pretrained("microsoft/layoutlmv3-base")191 192#def process_question(question, document):193# #print(f"Debug - Processing Question: {question}")194# 195# encoding = processor(document, question, return_tensors="pt")196# #print(f"Debug - Encoding Input IDs: {encoding.input_ids}")197#198# outputs = model(**encoding)199# #print(f"Debug - Model Outputs: {outputs}") 200#201# predicted_start_idx = outputs.start_logits.argmax(-1).item()202# predicted_end_idx = outputs.end_logits.argmax(-1).item()203#204# # Check if indices are valid205# if predicted_start_idx < 0 or predicted_end_idx < 0:206# print(f"Warning - Invalid prediction indices: start={predicted_start_idx}, end={predicted_end_idx}")207# return ""208#209# answer_tokens = encoding.input_ids.squeeze()[predicted_start_idx: predicted_end_idx + 1]210# answer = processor.tokenizer.decode(answer_tokens)211#212# return answer213 214# Older iteration of the code, retaining for emergencies ?215#def process_question(question, document):216# if not question or document is None:217# return None, None, None218#219# text_value = None220# predictions = run_pipeline(question, document)221#222# for i, p in enumerate(ensure_list(predictions)):223# if i == 0:224# text_value = p["answer"]225# else:226# # Keep the code around to produce multiple boxes, but only show the top227# # prediction for now228# break229# 230# return text_value231 232#def parse_ticket_image(image, question):233# """Basically just runs through these questions for the document"""234# # Processing the image 235# if image: 236# try: 237# if image.mode != "RGB":238# document = image.convert("RGB")239# else: 240# document = image241# except Exception as e:242# traceback.print_exc()243# error = str(e)244# 245# 246# # Define questions you want to ask the model247# 248# questions = [249# "What is the ticket number?", 250# "What is the type of grain (For example: corn, soybeans, wheat)?", 251# "What is the date?", 252# "What is the time?", 253# "What is the gross weight?", 254# "What is the tare weight?", 255# "What is the net weight?", 256# "What is the moisture (moist) percentage?", 257# "What is the damage percentage?",258# "What is the gross units?",259# "What is the dock units?", 260# "What is the comment?", 261# "What is the assembly number?",262# ]263# 264# # Use the model to answer each question265# answers = {}266# for q in questions: 267# print(f"Question: {q}")268# answer_text = process_question(q, document)269# print(f"Answer Text extracted here: {answer_text}")270# answers[q] = answer_text271# 272# 273# ticket_number = answers["What is the ticket number?"]274# grain_type = answers["What is the type of grain (For example: corn, soybeans, wheat)?"]275# date = answers["What is the date?"]276# time = answers["What is the time?"]277# gross_weight = answers["What is the gross weight?"]278# tare_weight = answers["What is the tare weight?"]279# net_weight = answers["What is the net weight?"]280# moisture = answers["What is the moisture (moist) percentage?"]281# damage = answers["What is the damage percentage?"]282# gross_units = answers["What is the gross units?"]283# dock_units = answers["What is the dock units?"]284# comment = answers["What is the comment?"]285# assembly_number = answers["What is the assembly number?"]286#287# 288# # Create a structured format (like a table) using pandas289# data = {290# "Ticket Number": [ticket_number],291# "Grain Type": [grain_type],292# "Assembly Number": [assembly_number],293# "Date": [date],294# "Time": [time],295# "Gross Weight": [gross_weight],296# "Tare Weight": [tare_weight],297# "Net Weight": [net_weight],298# "Moisture": [moisture],299# "Damage": [damage],300# "Gross Units": [gross_units],301# "Dock Units": [dock_units],302# "Comment": [comment],303# }304# df = pd.DataFrame(data)305# 306# return df307 308 309 310"""311For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface312"""313demo = gr.Interface(314 fn=parse_ticket_image, 315 inputs=[gr.Image(label= "Upload your Grain Scale Ticket", type="pil")],316 outputs=[gr.Dataframe(headers=["Field", "Value"], label="Extracted Grain Scale Ticket Data")],317)318 319 320if __name__ == "__main__":321 demo.launch()