import os import shutil from fuzzywuzzy import fuzz from tqdm import tqdm from PIL import Image import requests # from surya.layout import LayoutPredictor from doctr.io import DocumentFile from pdf2image import convert_from_path import pymupdf # from doctr.models import ocr_predictor import numpy as np from time import time pipe = None layout_predictor = None MAX_BLOCK_MATCHES = 2 MAX_LINE_MATCHES = 5 CUT_OFF_THRESHOLD = 60 QUESTION_WEIGHT = 0.2 ANSWER_WEIGHT = 0.8 LEVEL = "line" jpg_options = { "quality" : 100, "progressive": True, "optimize" : False } stop_words = {'what', 'is', 'the', 'this', 'that', 'these', 'those', 'which', 'how', 'why', 'where', 'when', 'who', 'will', 'be', 'and', 'or', 'in', 'at', 'to', 'for', 'of', 'with', 'by'} def longest_consecutive_range(indices): if not indices: return [] indices = sorted(set(indices)) longest = [] current = [indices[0]] for i in range(1, len(indices)): if indices[i] == indices[i - 1] + 1: current.append(indices[i]) else: if len(current) > len(longest): longest = current current = [indices[i]] if len(current) > len(longest): longest = current return longest def get_word_level_matches(answer_text, top_k_matches): bboxes = [] for match in top_k_matches: indices = [] for index, word in enumerate(match['words']): if word['text'].lower() in answer_text.lower(): # bboxes.append(word['bbox']) indices.append(index) longest_indices = longest_consecutive_range(indices) for index in longest_indices: bboxes.append(match['words'][index]['bbox']) return bboxes def get_matched_regions(question_text, target_text, predictions, level): question_terms = [word.lower() for word in question_text.split() if word.lower() not in stop_words] matched_regions = [] for region in predictions: region_text = region['text'] region_copy = region.copy() if target_text.lower() in region_text.lower(): region_copy['match_score'] = 100 region_copy['match_details'] = { 'exact_match': True, 'answer_score': 100, 'question_score': 100 } matched_regions.append(region_copy) continue partial_score = fuzz.partial_ratio(target_text.lower(), region_text.lower()) token_score = fuzz.token_set_ratio(target_text.lower(), region_text.lower()) # Calculate length factor (preference for longer matches that contain meaningful content) target_len = len(target_text) region_len = len(region_text) length_factor = min(1.0, region_len / min(50, target_len)) # Cap at 1.0, adapt based on target length # Combine scores for answer with weights # Higher weight to token matching for longer texts, higher weight to partial matching for shorter texts if region_len > 10: answer_score = (partial_score * 0.3) + (token_score * 0.5) + (length_factor * 100 * 0.2) else: # For very short texts, reduce their overall score unless they're exact matches answer_score = (partial_score * 0.3) + (token_score * 0.4) + (length_factor * 100 * 0.3) if region_len < 5 and partial_score < 100: answer_score *= 0.5 # Penalize very short inexact matches # penalize shorter region_texts if region_len < 5: answer_score *= 0.5 # Calculate fuzzy match scores for question terms using both methods partial_question_scores = [fuzz.partial_ratio(term, region_text.lower()) for term in question_terms] token_question_scores = [fuzz.token_set_ratio(term, region_text.lower()) for term in question_terms] # Get best scores for question terms best_partial_question = max(partial_question_scores) if partial_question_scores else 0 best_token_question = max(token_question_scores) if token_question_scores else 0 # Combine question scores question_score = (best_partial_question * 0.4) + (best_token_question * 0.6) # Combine scores (giving more weight to answer matches) combined_score = (answer_score * ANSWER_WEIGHT) + (question_score * QUESTION_WEIGHT) # print(combined_score) if combined_score >= CUT_OFF_THRESHOLD: region_copy['match_score'] = combined_score region_copy['match_details'] = { 'exact_match': False, 'answer_score': answer_score, 'question_score': question_score, 'answer_weight': ANSWER_WEIGHT, 'question_weight': QUESTION_WEIGHT } matched_regions.append(region_copy) matched_regions.sort(key=lambda x: x['match_score'], reverse=True) # If no matches, reduce threshold by 20 and get the topmost single output if not matched_regions: new_threshold = max(CUT_OFF_THRESHOLD - 20, 0) # Prevent negative threshold matched_regions = [region for region in matched_regions if region['match_score'] >= new_threshold] matched_regions.sort(key=lambda x: x['match_score'], reverse=True) if matched_regions: matched_regions = [matched_regions[0]] # Only keep the topmost single output if level == "block": top_matches = matched_regions[:MAX_BLOCK_MATCHES] elif level == "line": top_matches = matched_regions[:MAX_LINE_MATCHES] return top_matches def get_processed_text_for_llm(block_predictions, gap): final_text = "" for block_data in block_predictions: final_text += block_data['text'] + gap return final_text def get_page_number(block_bboxes): pages = {} for block in block_bboxes: if block['page'] not in pages: pages[block['page']] = 1 else: pages[block['page']] += 1 print(pages) max_page = max(pages, key=pages.get) return max_page def predict_output(document_path, question, pipe, layout_predictor, model, model_type, document_type="image"): predicted_answer = None block_box_predictions = None line_box_predictions = None word_box_predictions = None point_box_predictions = None curr_time = time() line_predictions, pages_count = get_line_predictions(document_path, model, document_type) line_time = time() print(f"Done with line predictions in {line_time - curr_time} seconds") curr_time = time() if(document_type == "pdf" and pages_count < 3): block_predictions = get_block_predictions(document_path, layout_predictor, model, document_type) gap = '\n\n\n' else: block_predictions = line_predictions gap = '\n' block_time = time() print(f"Done with block predictions in {block_time - line_time} seconds") # exit() # print(line_predictions) # print(block_predictions) curr_time = time() if model_type == "MGVG" or document_type=="pdf": processed_text_for_llm = get_processed_text_for_llm(block_predictions, gap=gap) # print("Processed Text for LLM: ", processed_text_for_llm) predicted_answer = generate_llm_answer(question, processed_text_for_llm, pipe) elif model_type == "IndoDocs": predicted_answer = generate_via_inhouse_model_answer(question, document_path) llm_time = time() print(f"Done with LLM in {llm_time - curr_time} seconds") print("LLM Answer: ", predicted_answer) total_algo_time = time() # print(predicted_answer) curr_time = time() line_matches = get_matched_regions(question, predicted_answer, line_predictions, "line") block_bboxes = get_matched_regions(question, predicted_answer, block_predictions, "block") match_time = time() print(f"Done with match in {match_time - curr_time} seconds") if document_type == "pdf": current_page = get_page_number(block_bboxes) else: current_page = -1 if(current_page != -1): predicted_answer = "Answer predicted from page: " + str(current_page+1) + "\n" + predicted_answer block_box_predictions = [] for match in block_bboxes: block_box_predictions.append(match['bbox']) line_box_predictions = [] for match in line_matches: # print(match['page'], match['bbox']) if current_page == -1 or match['page'] == current_page: line_box_predictions.append(match['bbox']) # for line in line_box_predictions: # print(line) curr_time = time() word_box_predictions = get_word_level_matches(predicted_answer, top_k_matches=line_matches) word_time = time() print(f"Done with word in {word_time - curr_time} seconds") curr_time = time() point_box_predictions = get_point_level_matches(block_box_predictions, line_box_predictions, word_box_predictions) point_time = time() print(f"Done with point in {point_time - curr_time} seconds") print(f"Total algo time: {time() - total_algo_time} seconds") # print(block_box_predictions) # print(line_box_predictions) # print(word_box_predictions) # print(point_box_predictions) # print(predicted_answer) return predicted_answer, block_box_predictions, line_box_predictions, word_box_predictions, point_box_predictions, current_page def calculate_midpoint_of_bboxes(bboxes): if not bboxes: return None # Convert to numpy array for easier manipulation bboxes = np.array(bboxes) # Find the extreme points of all bboxes combined min_x = np.min(bboxes[:, 0]) min_y = np.min(bboxes[:, 1]) max_x = np.max(bboxes[:, 2]) max_y = np.max(bboxes[:, 3]) # Calculate midpoint midpoint_x = (min_x + max_x) / 2 midpoint_y = (min_y + max_y) / 2 return round(midpoint_x, 2), round(midpoint_y, 2) def get_point_level_matches(block_box_predictions, line_box_predictions, word_box_predictions): point_box_predictions = [] if len(block_box_predictions) ==1: try: x, y = calculate_midpoint_of_bboxes(block_box_predictions) point_box_predictions = [[x, y]] # print(x, y) except: try: x, y = calculate_midpoint_of_bboxes(line_box_predictions) point_box_predictions = [[x, y]] except: point_box_predictions = [] else: points = [] for block_bbox in block_box_predictions: try: x, y = calculate_midpoint_of_bboxes(block_bbox) points.append([x, y]) except: continue point_box_predictions = points return point_box_predictions def generate_via_inhouse_model_answer(question, image_path, api_key="VISION-TEAM", max_tokens=512, temperature=0.7, endpoint="http://103.207.148.38:9000/api/v1/chat/upload"): headers = { "x-api-key": api_key # or whatever the Swagger UI says } files = { "image": open(image_path, "rb") } data = { "text": question, "max_tokens": str(max_tokens), "temperature": str(temperature) } try: response = requests.post(endpoint, headers=headers, files=files, data=data) response.raise_for_status() result = response.json() except requests.exceptions.RequestException as e: return {"error": str(e)} return result['response']['choices'][0]['message']['content'] def generate_llm_answer(question, context, pipe): prompt = f"""You are given a question and context. Your task is to find and return the best possible answer to the question using only the context as it is. Do not generate summaries, paraphrased content, or any additional explanations including any preamble and postamble. Return only the exact phrase or sentence fragment from the context that answers the question. If the answer is not found in the context, return: Answer not found in context. Question: {question} Context: {context} Answer: """ messages = [ {"role": "user", "content": prompt}] result = pipe(messages, max_new_tokens=512, do_sample=True, temperature=0.7) ans = result[0]["generated_text"][1]['content'] return ans def get_line_predictions(document_path, model, document_type): current_dir = os.getcwd() if document_type == "pdf": output_file = simple_counter_generator("page", ".jpg") current_dir = os.getcwd() temp_output_folder = os.path.join(current_dir, "temp_output_folder/") # delete the temp_output_folder if os.path.exists(temp_output_folder): shutil.rmtree(temp_output_folder) if not os.path.exists(temp_output_folder): os.makedirs(temp_output_folder) # output_file = simple_counter_generator("page", ".jpg") # convert_from_path(document_path, output_folder=temp_output_folder, dpi=300, fmt='jpeg', jpegopt= jpg_options, output_file=output_file) doc = pymupdf.open(document_path) # open document for page in doc: # iterate through the pages pix = page.get_pixmap() # render page to an image pix.save(f"{temp_output_folder}/{page.number}.png") images_path = sorted(os.listdir(temp_output_folder)) else: images_path = [os.path.join(current_dir, document_path)] print(images_path) block_predictions = [] # print(document_path) # if document_type == "pdf": # doc = DocumentFile.from_pdf(document_path) # else: # doc = DocumentFile.from_images(document_path) # result = model(doc) line_predictions = [] pages_count = -1 for image_path in images_path: pages_count += 1 if(len(images_path) > 1): doc = DocumentFile.from_images(os.path.join(temp_output_folder, image_path)) else: doc = DocumentFile.from_images(image_path) result = model(doc) for page in result.pages: dim = tuple(reversed(page.dimensions)) for block in page.blocks: for line in block.lines: output = {} geo = line.geometry a = list(a*b for a,b in zip(geo[0],dim)) b = list(a*b for a,b in zip(geo[1],dim)) x1 = round(a[0], 2).astype(float) y1 = round(a[1], 2).astype(float) x2 = round(b[0], 2).astype(float) y2 = round(b[1], 2).astype(float) line_bbox = [x1, y1, x2, y2] sent = [] words_data = [] for word in line.words: word_data = {} sent.append(word.value) geo = word.geometry a = list(a*b for a,b in zip(geo[0],dim)) b = list(a*b for a,b in zip(geo[1],dim)) x1 = round(a[0], 2).astype(float) y1 = round(a[1], 2).astype(float) x2 = round(b[0], 2).astype(float) y2 = round(b[1], 2).astype(float) bbox = [x1, y1, x2, y2] word_data['bbox'] = bbox word_data['text'] = word.value words_data.append(word_data) output['bbox'] = line_bbox output['text'] = " ".join(sent) output['words'] = words_data output['page'] = pages_count line_predictions.append(output) return line_predictions, pages_count def get_block_predictions(document_path, layout_predictor, model, document_type): current_dir = os.getcwd() if document_type == "pdf": output_file = simple_counter_generator("page", ".jpg") current_dir = os.getcwd() temp_output_folder = os.path.join(current_dir, "temp_output_folder/") # delete the temp_output_folder if os.path.exists(temp_output_folder): shutil.rmtree(temp_output_folder) if not os.path.exists(temp_output_folder): os.makedirs(temp_output_folder) # output_file = simple_counter_generator("page", ".jpg") # convert_from_path(document_path, output_folder=temp_output_folder, dpi=300, fmt='jpeg', jpegopt= jpg_options, output_file=output_file) doc = pymupdf.open(document_path) # open document for page in doc: # iterate through the pages pix = page.get_pixmap() # render page to an image pix.save(f"{temp_output_folder}/{page.number}.png") images_path = sorted(os.listdir(temp_output_folder)) else: images_path = [os.path.join(current_dir, document_path)] # print(images_path) block_predictions = [] page_count = -1 for image_path in images_path: page_count += 1 if(len(images_path) > 1): image = Image.open(os.path.join(temp_output_folder, image_path)) else: image = Image.open(os.path.join(current_dir, document_path)) # print(image_path) # print(image) layout_predictions = layout_predictor([image]) for block in layout_predictions[0].bboxes: output = {} bbox = [int(x) for x in block.bbox] cropped_image = image.crop(bbox) cropped_image.save(f'temp.png') doc = DocumentFile.from_images('temp.png') result = model(doc) text = [] for page in result.pages: for block in page.blocks: for line in block.lines: for word in line.words: text.append(word.value) output['bbox'] = bbox output['text'] = " ".join(text) output['page'] = page_count block_predictions.append(output) return block_predictions def simple_counter_generator(prefix="", suffix=""): while True: yield 'p' # from doctr.models import ocr_predictor # model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True) # # from transformers import pipeline # # def load_llm_model(device): # # pipe = pipeline("text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device=device) # # return pipe # # pipe = load_llm_model("cuda") # pipe = None # # from surya.layout import LayoutPredictor # # layout_predictor = LayoutPredictor() # layout_predictor = None # document_path = "sample.pdf" # question = "What is the subject of the circular?" # answer, block_box_predictions, line_box_predictions, word_box_predictions, point_box_predictions = predict_output(document_path, question, pipe, layout_predictor, model, "Inhouse", document_type="pdf") # print(answer) # print(block_box_predictions) # print(line_box_predictions) # print(word_box_predictions) # print(point_box_predictions)