import openai import os, sys, json, inspect, glob, tiktoken, shutil, yaml, torch, logging import openpyxl from openpyxl import Workbook, load_workbook import google.generativeai as genai import vertexai from langchain_openai import AzureChatOpenAI from transformers import TrOCRProcessor, VisionEncoderDecoderModel currentdir = os.path.dirname(os.path.abspath( inspect.getfile(inspect.currentframe()))) parentdir = os.path.dirname(currentdir) sys.path.append(parentdir) parentdir = os.path.dirname(parentdir) sys.path.append(parentdir) from general_utils import get_cfg_from_full_path # from embeddings_db import VoucherVisionEmbedding from OCR_google_cloud_vision import OCRGoogle from vouchervision.LLM_OpenAI import OpenAIHandler from vouchervision.LLM_GooglePalm2 import GooglePalm2Handler from vouchervision.LLM_GoogleGemini import GoogleGeminiHandler from vouchervision.LLM_MistralAI import MistralHandler from vouchervision.LLM_local_cpu_MistralAI import LocalCPUMistralHandler from vouchervision.LLM_local_MistralAI import LocalMistralHandler #call_llm_local_MistralAI_8x7b from vouchervision.utils_LLM import remove_colons_and_double_apostrophes # from LLM_Falcon import OCR_to_dict_Falcon from prompt_catalog import PromptCatalog from vouchervision.model_maps import ModelMaps ''' * For the prefix_removal, the image names have 'MICH-V-' prior to the barcode, so that is used for matching but removed for output. * There is also code active to replace the LLM-predicted "Catalog Number" with the correct number since it is known. The LLMs to usually assign the barcode to the correct field, but it's not needed since it is already known. - Look for ####################### Catalog Number pre-defined ''' class VoucherVision(): def __init__(self, cfg, logger, dir_home, path_custom_prompts, Project, Dirs): self.cfg = cfg self.logger = logger self.dir_home = dir_home self.path_custom_prompts = path_custom_prompts self.Project = Project self.Dirs = Dirs self.headers = None self.prompt_version = None # self.trOCR_model_version = "microsoft/trocr-large-handwritten" self.trOCR_model_version = "microsoft/trocr-base-handwritten" self.trOCR_processor = None self.trOCR_model = None self.set_API_keys() self.setup() def setup(self): self.logger.name = f'[Transcription]' self.logger.info(f'Setting up OCR and LLM') self.db_name = self.cfg['leafmachine']['project']['embeddings_database_name'] self.path_domain_knowledge = self.cfg['leafmachine']['project']['path_to_domain_knowledge_xlsx'] self.build_new_db = self.cfg['leafmachine']['project']['build_new_embeddings_database'] self.continue_run_from_partial_xlsx = self.cfg['leafmachine']['project']['continue_run_from_partial_xlsx'] self.prefix_removal = self.cfg['leafmachine']['project']['prefix_removal'] self.suffix_removal = self.cfg['leafmachine']['project']['suffix_removal'] self.catalog_numerical_only = self.cfg['leafmachine']['project']['catalog_numerical_only'] self.prompt_version0 = self.cfg['leafmachine']['project']['prompt_version'] self.use_domain_knowledge = self.cfg['leafmachine']['project']['use_domain_knowledge'] self.catalog_name_options = ["Catalog Number", "catalog_number"] self.utility_headers = ["filename", "WFO_override_OCR", "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement", "GEO_override_OCR", "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat", "GEO_decimal_long","GEO_city", "GEO_county", "GEO_state", "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent", "tokens_in", "tokens_out", "path_to_crop","path_to_original","path_to_content","path_to_helper",] self.do_create_OCR_helper_image = self.cfg['leafmachine']['do_create_OCR_helper_image'] self.map_prompt_versions() self.map_dir_labels() self.map_API_options() # self.init_embeddings() self.init_transcription_xlsx() self.init_trOCR_model() '''Logging''' self.logger.info(f'Transcribing dataset --- {self.dir_labels}') self.logger.info(f'Saving transcription batch to --- {self.path_transcription}') self.logger.info(f'Saving individual transcription files to --- {self.Dirs.transcription_ind}') self.logger.info(f'Starting transcription...') self.logger.info(f' LLM MODEL --> {self.version_name}') self.logger.info(f' Using Azure API --> {self.is_azure}') self.logger.info(f' Model name passed to API --> {self.model_name}') self.logger.info(f' API access token is found in PRIVATE_DATA.yaml --> {self.has_key}') def init_trOCR_model(self): lgr = logging.getLogger('transformers') lgr.setLevel(logging.ERROR) self.trOCR_processor = TrOCRProcessor.from_pretrained(self.trOCR_model_version) self.trOCR_model = VisionEncoderDecoderModel.from_pretrained(self.trOCR_model_version) # Check for GPU availability self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.trOCR_model.to(self.device) def map_API_options(self): self.chat_version = self.cfg['leafmachine']['LLM_version'] # Get the required values from ModelMaps self.model_name = ModelMaps.get_version_mapping_cost(self.chat_version) self.is_azure = ModelMaps.get_version_mapping_is_azure(self.chat_version) self.has_key = ModelMaps.get_version_has_key(self.chat_version, self.has_key_openai, self.has_key_azure_openai, self.has_key_palm2, self.has_key_mistral) # Check if the version is supported if self.model_name is None: supported_LLMs = ", ".join(ModelMaps.get_models_gui_list()) raise Exception(f"Unsupported LLM: {self.chat_version}. Requires one of: {supported_LLMs}") self.version_name = self.chat_version def map_prompt_versions(self): self.prompt_version_map = { "Version 1": "prompt_v1_verbose", "Version 1 No Domain Knowledge": "prompt_v1_verbose_noDomainKnowledge", "Version 2": "prompt_v2_json_rules", "Version 1 PaLM 2": 'prompt_v1_palm2', "Version 1 PaLM 2 No Domain Knowledge": 'prompt_v1_palm2_noDomainKnowledge', "Version 2 PaLM 2": 'prompt_v2_palm2', } self.prompt_version = self.prompt_version_map.get(self.prompt_version0, self.path_custom_prompts) self.is_predefined_prompt = self.is_in_prompt_version_map(self.prompt_version) def is_in_prompt_version_map(self, value): return value in self.prompt_version_map.values() # def init_embeddings(self): # if self.use_domain_knowledge: # self.logger.info(f'*** USING DOMAIN KNOWLEDGE ***') # self.logger.info(f'*** Initializing vector embeddings database ***') # self.initialize_embeddings() # else: # self.Voucher_Vision_Embedding = None def map_dir_labels(self): if self.cfg['leafmachine']['use_RGB_label_images']: self.dir_labels = os.path.join(self.Dirs.save_per_annotation_class,'label') else: self.dir_labels = self.Dirs.save_original # Use glob to get all image paths in the directory self.img_paths = glob.glob(os.path.join(self.dir_labels, "*")) def load_rules_config(self): with open(self.path_custom_prompts, 'r') as stream: try: return yaml.safe_load(stream) except yaml.YAMLError as exc: print(exc) return None def generate_xlsx_headers(self): # Extract headers from the 'Dictionary' keys in the JSON template rules # xlsx_headers = list(self.rules_config_json['rules']["Dictionary"].keys()) xlsx_headers = list(self.rules_config_json['rules'].keys()) xlsx_headers = xlsx_headers + self.utility_headers return xlsx_headers def init_transcription_xlsx(self): self.HEADERS_v1_n22 = ["Catalog Number","Genus","Species","subspecies","variety","forma","Country","State","County","Locality Name","Min Elevation","Max Elevation","Elevation Units","Verbatim Coordinates","Datum","Cultivated","Habitat","Collectors","Collector Number","Verbatim Date","Date","End Date"] self.HEADERS_v2_n26 = ["catalog_number","genus","species","subspecies","variety","forma","country","state","county","locality_name","min_elevation","max_elevation","elevation_units","verbatim_coordinates","decimal_coordinates","datum","cultivated","habitat","plant_description","collectors","collector_number","determined_by","multiple_names","verbatim_date","date","end_date"] self.HEADERS_v1_n22 = self.HEADERS_v1_n22 + self.utility_headers self.HEADERS_v2_n26 = self.HEADERS_v2_n26 + self.utility_headers # Initialize output file self.path_transcription = os.path.join(self.Dirs.transcription,"transcribed.xlsx") if self.prompt_version in ['prompt_v2_json_rules','prompt_v2_palm2']: self.headers = self.HEADERS_v2_n26 self.headers_used = 'HEADERS_v2_n26' elif self.prompt_version in ['prompt_v1_verbose', 'prompt_v1_verbose_noDomainKnowledge','prompt_v1_palm2', 'prompt_v1_palm2_noDomainKnowledge']: self.headers = self.HEADERS_v1_n22 self.headers_used = 'HEADERS_v1_n22' else: if not self.is_predefined_prompt: # Load the rules configuration self.rules_config_json = self.load_rules_config() # Generate the headers from the configuration self.headers = self.generate_xlsx_headers() # Set the headers used to the dynamically generated headers self.headers_used = 'CUSTOM' else: # If it's a predefined prompt, raise an exception as we don't have further instructions raise ValueError("Predefined prompt is not handled in this context.") self.create_or_load_excel_with_headers(os.path.join(self.Dirs.transcription,"transcribed.xlsx"), self.headers) def create_or_load_excel_with_headers(self, file_path, headers, show_head=False): output_dir_names = ['Archival_Components', 'Config_File', 'Cropped_Images', 'Logs', 'Original_Images', 'Transcription'] self.completed_specimens = [] # Check if the file exists and it's not None if self.continue_run_from_partial_xlsx is not None and os.path.isfile(self.continue_run_from_partial_xlsx): workbook = load_workbook(filename=self.continue_run_from_partial_xlsx) sheet = workbook.active show_head=True # Identify the 'path_to_crop' column try: path_to_crop_col = headers.index('path_to_crop') + 1 path_to_original_col = headers.index('path_to_original') + 1 path_to_content_col = headers.index('path_to_content') + 1 path_to_helper_col = headers.index('path_to_helper') + 1 # self.completed_specimens = list(sheet.iter_cols(min_col=path_to_crop_col, max_col=path_to_crop_col, values_only=True, min_row=2)) except ValueError: print("'path_to_crop' not found in the header row.") path_to_crop = list(sheet.iter_cols(min_col=path_to_crop_col, max_col=path_to_crop_col, values_only=True, min_row=2)) path_to_original = list(sheet.iter_cols(min_col=path_to_original_col, max_col=path_to_original_col, values_only=True, min_row=2)) path_to_content = list(sheet.iter_cols(min_col=path_to_content_col, max_col=path_to_content_col, values_only=True, min_row=2)) path_to_helper = list(sheet.iter_cols(min_col=path_to_helper_col, max_col=path_to_helper_col, values_only=True, min_row=2)) others = [path_to_crop_col, path_to_original_col, path_to_content_col, path_to_helper_col] jsons = [path_to_content_col, path_to_helper_col] for cell in path_to_crop[0]: old_path = cell new_path = file_path for dir_name in output_dir_names: if dir_name in old_path: old_path_parts = old_path.split(dir_name) new_path_parts = new_path.split('Transcription') updated_path = new_path_parts[0] + dir_name + old_path_parts[1] self.completed_specimens.append(os.path.basename(updated_path)) print(f"{len(self.completed_specimens)} images are already completed") ### Copy the JSON files over for colu in jsons: cell = next(sheet.iter_rows(min_row=2, min_col=colu, max_col=colu))[0] old_path = cell.value new_path = file_path old_path_parts = old_path.split('Transcription') new_path_parts = new_path.split('Transcription') updated_path = new_path_parts[0] + 'Transcription' + old_path_parts[1] # Copy files old_dir = os.path.dirname(old_path) new_dir = os.path.dirname(updated_path) # Check if old_dir exists and it's a directory if os.path.exists(old_dir) and os.path.isdir(old_dir): # Check if new_dir exists. If not, create it. if not os.path.exists(new_dir): os.makedirs(new_dir) # Iterate through all files in old_dir and copy each to new_dir for filename in os.listdir(old_dir): shutil.copy2(os.path.join(old_dir, filename), new_dir) # copy2 preserves metadata ### Update the file names for colu in others: for row in sheet.iter_rows(min_row=2, min_col=colu, max_col=colu): for cell in row: old_path = cell.value new_path = file_path for dir_name in output_dir_names: if dir_name in old_path: old_path_parts = old_path.split(dir_name) new_path_parts = new_path.split('Transcription') updated_path = new_path_parts[0] + dir_name + old_path_parts[1] cell.value = updated_path show_head=True else: # Create a new workbook and select the active worksheet workbook = Workbook() sheet = workbook.active # Write headers in the first row for i, header in enumerate(headers, start=1): sheet.cell(row=1, column=i, value=header) self.completed_specimens = [] # Save the workbook workbook.save(file_path) if show_head: print("continue_run_from_partial_xlsx:") for i, row in enumerate(sheet.iter_rows(values_only=True)): print(row) if i == 3: # print the first 5 rows (0-indexed) print("\n") break def add_data_to_excel_from_response(self, path_transcription, response, WFO_record, GEO_record, filename_without_extension, path_to_crop, path_to_content, path_to_helper, nt_in, nt_out): geo_headers = ["GEO_override_OCR", "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat", "GEO_decimal_long","GEO_city", "GEO_county", "GEO_state", "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent",] # WFO_candidate_names is separate, bc it may be type --> list wfo_headers = ["WFO_override_OCR", "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_placement"] wb = openpyxl.load_workbook(path_transcription) sheet = wb.active # find the next empty row next_row = sheet.max_row + 1 if isinstance(response, str): try: response = json.loads(response) except json.JSONDecodeError: print(f"Failed to parse response: {response}") return # iterate over headers in the first row for i, header in enumerate(sheet[1], start=1): # check if header value is in response keys if (header.value in response) and (header.value not in self.catalog_name_options): ####################### Catalog Number pre-defined # check if the response value is a dictionary if isinstance(response[header.value], dict): # if it is a dictionary, extract the 'value' field cell_value = response[header.value].get('value', '') else: # if it's not a dictionary, use it directly cell_value = response[header.value] try: # write the value to the cell sheet.cell(row=next_row, column=i, value=cell_value) except: sheet.cell(row=next_row, column=i, value=cell_value[0]) elif header.value in self.catalog_name_options: # if self.prefix_removal: # filename_without_extension = filename_without_extension.replace(self.prefix_removal, "") # if self.suffix_removal: # filename_without_extension = filename_without_extension.replace(self.suffix_removal, "") # if self.catalog_numerical_only: # filename_without_extension = self.remove_non_numbers(filename_without_extension) sheet.cell(row=next_row, column=i, value=filename_without_extension) elif header.value == "path_to_crop": sheet.cell(row=next_row, column=i, value=path_to_crop) elif header.value == "path_to_original": if self.cfg['leafmachine']['use_RGB_label_images']: fname = os.path.basename(path_to_crop) base = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(path_to_crop)))) path_to_original = os.path.join(base, 'Original_Images', fname) sheet.cell(row=next_row, column=i, value=path_to_original) else: fname = os.path.basename(path_to_crop) base = os.path.dirname(os.path.dirname(path_to_crop)) path_to_original = os.path.join(base, 'Original_Images', fname) sheet.cell(row=next_row, column=i, value=path_to_original) elif header.value == "path_to_content": sheet.cell(row=next_row, column=i, value=path_to_content) elif header.value == "path_to_helper": sheet.cell(row=next_row, column=i, value=path_to_helper) elif header.value == "tokens_in": sheet.cell(row=next_row, column=i, value=nt_in) elif header.value == "tokens_out": sheet.cell(row=next_row, column=i, value=nt_out) elif header.value == "filename": sheet.cell(row=next_row, column=i, value=filename_without_extension) # "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement" elif header.value in wfo_headers: sheet.cell(row=next_row, column=i, value=WFO_record.get(header.value, '')) # elif header.value == "WFO_exact_match": # sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_exact_match",'')) # elif header.value == "WFO_exact_match_name": # sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_exact_match_name",'')) # elif header.value == "WFO_best_match": # sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_best_match",'')) # elif header.value == "WFO_placement": # sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_placement",'')) elif header.value == "WFO_candidate_names": candidate_names = WFO_record.get("WFO_candidate_names", '') # Check if candidate_names is a list and convert to a string if it is if isinstance(candidate_names, list): candidate_names_str = '|'.join(candidate_names) else: candidate_names_str = candidate_names sheet.cell(row=next_row, column=i, value=candidate_names_str) # "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat", "GEO_decimal_long", # "GEO_city", "GEO_county", "GEO_state", "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent" elif header.value in geo_headers: sheet.cell(row=next_row, column=i, value=GEO_record.get(header.value, '')) # save the workbook wb.save(path_transcription) def has_API_key(self, val): if val != '': return True else: return False def set_API_keys(self): self.dir_home = os.path.dirname(os.path.dirname(__file__)) self.path_cfg_private = os.path.join(self.dir_home, 'PRIVATE_DATA.yaml') self.cfg_private = get_cfg_from_full_path(self.path_cfg_private) self.has_key_openai = self.has_API_key(self.cfg_private['openai']['OPENAI_API_KEY']) self.has_key_azure_openai = self.has_API_key(self.cfg_private['openai_azure']['api_version']) self.has_key_palm2 = self.has_API_key(self.cfg_private['google_palm']['google_palm_api']) self.has_key_google_OCR = self.has_API_key(self.cfg_private['google_cloud']['path_json_file']) self.has_key_mistral = self.has_API_key(self.cfg_private['mistral']['mistral_key']) self.has_key_here = self.has_API_key(self.cfg_private['here']['api_key']) self.has_open_cage_geocode = self.has_API_key(self.cfg_private['open_cage_geocode']['api_key']) if self.has_key_openai: openai.api_key = self.cfg_private['openai']['OPENAI_API_KEY'] os.environ["OPENAI_API_KEY"] = self.cfg_private['openai']['OPENAI_API_KEY'] if self.has_key_azure_openai: # os.environ["OPENAI_API_KEY"] = self.cfg_private['openai_azure']['openai_api_key'] self.llm = AzureChatOpenAI( deployment_name = 'gpt-35-turbo',#'gpt-35-turbo', openai_api_version = self.cfg_private['openai_azure']['api_version'], openai_api_key = self.cfg_private['openai_azure']['openai_api_key'], azure_endpoint = self.cfg_private['openai_azure']['openai_api_base'], # openai_api_base=self.cfg_private['openai_azure']['openai_api_base'], openai_organization = self.cfg_private['openai_azure']['openai_organization'], # openai_api_type = self.cfg_private['openai_azure']['openai_api_type'] ) # Enable this for all LLMs EXCEPT GOOGLE name_check = self.cfg['leafmachine']['LLM_version'].lower().split(' ') if 'google' in name_check: pass else: if self.has_key_google_OCR: if os.path.exists(self.cfg_private['google_cloud']['path_json_file']): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.cfg_private['google_cloud']['path_json_file'] elif os.path.exists(self.cfg_private['google_cloud']['path_json_file_service_account2']): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.cfg_private['google_cloud']['path_json_file_service_account2'] else: raise f"Google JSON API key file not found" if self.has_key_palm2: os.environ['PALM'] = self.cfg_private['google_palm']['google_palm_api'] os.environ['PALM_PROJECT_ID'] = self.cfg_private['google_palm']['project_id'] os.environ['PALM_LOCATION'] = self.cfg_private['google_palm']['location'] os.environ['GOOGLE_API_KEY'] = self.cfg_private['google_palm']['google_palm_api'] # os.environ["GOOGLE_SERVICE_ACCOUNT"] = self.cfg_private['google_cloud']['path_json_file_service_account'] # vertexai.init(project='directed-curve-401601', location='us-central1') # genai.configure(api_key=os.environ['PALM']) if self.has_key_palm2: os.environ['MISTRAL_API_KEY'] = self.cfg_private['mistral']['mistral_key'] if self.has_key_here: os.environ['here_app_id'] = self.cfg_private['here']['app_id'] os.environ['here_api_key'] = self.cfg_private['here']['api_key'] if self.has_open_cage_geocode: os.environ['open_cage_geocode'] = self.cfg_private['open_cage_geocode']['api_key'] # def initialize_embeddings(self): # '''Loading embedding search __init__(self, db_name, path_domain_knowledge, logger, build_new_db=False, model_name="hkunlp/instructor-xl", device="cuda")''' # self.Voucher_Vision_Embedding = VoucherVisionEmbedding(self.db_name, self.path_domain_knowledge, logger=self.logger, build_new_db=self.build_new_db) def clean_catalog_number(self, data, filename_without_extension): #Cleans up the catalog number in data if it's a dict def modify_catalog_key(catalog_key, filename_without_extension, data): # Helper function to apply modifications on catalog number if catalog_key not in data: new_data = {catalog_key: None} data = {**new_data, **data} if self.prefix_removal: filename_without_extension = filename_without_extension.replace(self.prefix_removal, "") if self.suffix_removal: filename_without_extension = filename_without_extension.replace(self.suffix_removal, "") if self.catalog_numerical_only: filename_without_extension = self.remove_non_numbers(data[catalog_key]) data[catalog_key] = filename_without_extension return data if isinstance(data, dict): if self.headers_used == 'HEADERS_v1_n22': return modify_catalog_key("Catalog Number", filename_without_extension, data) elif self.headers_used in ['HEADERS_v2_n26', 'CUSTOM']: return modify_catalog_key("filename", filename_without_extension, data) else: raise ValueError("Invalid headers used.") else: raise TypeError("Data is not of type dict.") def write_json_to_file(self, filepath, data): '''Writes dictionary data to a JSON file.''' with open(filepath, 'w') as txt_file: if isinstance(data, dict): data = json.dumps(data, indent=4, sort_keys=False) txt_file.write(data) # def create_null_json(self): # return {} def remove_non_numbers(self, s): return ''.join([char for char in s if char.isdigit()]) def create_null_row(self, filename_without_extension, path_to_crop, path_to_content, path_to_helper): json_dict = {header: '' for header in self.headers} for header, value in json_dict.items(): if header == "path_to_crop": json_dict[header] = path_to_crop elif header == "path_to_original": fname = os.path.basename(path_to_crop) base = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(path_to_crop)))) path_to_original = os.path.join(base, 'Original_Images', fname) json_dict[header] = path_to_original elif header == "path_to_content": json_dict[header] = path_to_content elif header == "path_to_helper": json_dict[header] = path_to_helper elif header == "filename": json_dict[header] = filename_without_extension # "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement" elif header == "WFO_exact_match": json_dict[header] ='' elif header == "WFO_exact_match_name": json_dict[header] = '' elif header == "WFO_best_match": json_dict[header] = '' elif header == "WFO_candidate_names": json_dict[header] = '' elif header == "WFO_placement": json_dict[header] = '' return json_dict ################################################################################################################################## ################################################## OCR ################################################################## ################################################################################################################################## def perform_OCR_and_save_results(self, image_index, jpg_file_path_OCR_helper, txt_file_path_OCR, txt_file_path_OCR_bounds): self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Starting OCR') # self.OCR - None ### Process_image() runs the OCR for text, handwriting, trOCR AND creates the overlay image ocr_google = OCRGoogle(self.path_to_crop, self.cfg, self.trOCR_model_version, self.trOCR_model, self.trOCR_processor, self.device) ocr_google.process_image(self.do_create_OCR_helper_image, self.logger) self.OCR = ocr_google.OCR self.write_json_to_file(txt_file_path_OCR, ocr_google.OCR_JSON_to_file) self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Finished OCR') if len(self.OCR) > 0: ocr_google.overlay_image.save(jpg_file_path_OCR_helper) OCR_bounds = {} if ocr_google.hand_text_to_box_mapping is not None: OCR_bounds['OCR_bounds_handwritten'] = ocr_google.hand_text_to_box_mapping if ocr_google.normal_text_to_box_mapping is not None: OCR_bounds['OCR_bounds_printed'] = ocr_google.normal_text_to_box_mapping if ocr_google.trOCR_text_to_box_mapping is not None: OCR_bounds['OCR_bounds_trOCR'] = ocr_google.trOCR_text_to_box_mapping self.write_json_to_file(txt_file_path_OCR_bounds, OCR_bounds) self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Saved OCR Overlay Image') else: pass ########################################################################################################################### fix logic for no OCR ################################################################################################################################## ####################################################### LLM Switchboard ######################################################## ################################################################################################################################## def send_to_LLM(self, is_azure, progress_report, json_report, model_name): self.n_failed_LLM_calls = 0 self.n_failed_OCR = 0 final_JSON_response = None final_WFO_record = None final_GEO_record = None self.initialize_token_counters() self.update_progress_report_initial(progress_report) MODEL_NAME_FORMATTED = ModelMaps.get_API_name(model_name) name_parts = model_name.split("_") self.setup_JSON_dict_structure() llm_model = self.initialize_llm_model(self.logger, MODEL_NAME_FORMATTED, self.JSON_dict_structure, name_parts, is_azure, self.llm) prompts = [] for i, path_to_crop in enumerate(self.img_paths): self.update_progress_report_batch(progress_report, i) if self.should_skip_specimen(path_to_crop): self.log_skipping_specimen(path_to_crop) continue paths = self.generate_paths(path_to_crop, i) self.path_to_crop = path_to_crop filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper = paths self.perform_OCR_and_save_results(i, jpg_file_path_OCR_helper, txt_file_path_OCR, txt_file_path_OCR_bounds) if not self.OCR: self.n_failed_OCR += 1 else: # Format prompt prompt = self.setup_prompt() prompt = remove_colons_and_double_apostrophes(prompt) prompts.append(prompt) # Process prompts in batch if 'LOCAL' in name_parts and ('MISTRAL' in name_parts or 'MIXTRAL' in name_parts): batch_results = llm_model.call_llm_local_MistralAI(prompts) # Assuming this method is updated to handle batch processing # Process each result from the batch for i, result in enumerate(batch_results): response_candidate, nt_in, nt_out, WFO_record, GEO_record = result.values() self.n_failed_LLM_calls += 1 if response_candidate is None else 0 # Estimate n tokens returned self.logger.info(f'Prompt tokens IN --- {nt_in}') self.logger.info(f'Prompt tokens OUT --- {nt_out}') self.update_token_counters(nt_in, nt_out) final_JSON_response, final_WFO_record, final_GEO_record = self.update_final_response(response_candidate, WFO_record, GEO_record, paths, path_to_crop, nt_in, nt_out) self.log_completion_info(final_JSON_response) json_report.set_JSON(final_JSON_response, final_WFO_record, final_GEO_record) self.update_progress_report_final(progress_report) final_JSON_response = self.parse_final_json_response(final_JSON_response) return final_JSON_response, final_WFO_record, final_GEO_record, self.total_tokens_in, self.total_tokens_out ################################################################################################################################## ################################################## LLM Helper Funcs ############################################################## ################################################################################################################################## def initialize_llm_model(self, logger, model_name, JSON_dict_structure, name_parts, is_azure=None, llm_object=None): if 'LOCAL'in name_parts: if ('MIXTRAL' in name_parts) or ('MISTRAL' in name_parts): if 'CPU' in name_parts: return LocalCPUMistralHandler(logger, model_name, JSON_dict_structure) else: return LocalMistralHandler(logger, model_name, JSON_dict_structure) else: if 'PALM2' in name_parts: return GooglePalm2Handler(logger, model_name, JSON_dict_structure) elif 'GEMINI' in name_parts: return GoogleGeminiHandler(logger, model_name, JSON_dict_structure) elif 'MISTRAL' in name_parts and ('LOCAL' not in name_parts): return MistralHandler(logger, model_name, JSON_dict_structure) else: return OpenAIHandler(logger, model_name, JSON_dict_structure, is_azure, llm_object) def setup_prompt(self): Catalog = PromptCatalog() prompt, _ = Catalog.prompt_SLTP(self.path_custom_prompts, OCR=self.OCR) return prompt def setup_JSON_dict_structure(self): Catalog = PromptCatalog() _, self.JSON_dict_structure = Catalog.prompt_SLTP(self.path_custom_prompts, OCR='Text') def initialize_token_counters(self): self.total_tokens_in = 0 self.total_tokens_out = 0 def update_progress_report_initial(self, progress_report): if progress_report is not None: progress_report.set_n_batches(len(self.img_paths)) def update_progress_report_batch(self, progress_report, batch_index): if progress_report is not None: progress_report.update_batch(f"Working on image {batch_index + 1} of {len(self.img_paths)}") def should_skip_specimen(self, path_to_crop): return os.path.basename(path_to_crop) in self.completed_specimens def log_skipping_specimen(self, path_to_crop): self.logger.info(f'[Skipping] specimen {os.path.basename(path_to_crop)} already processed') def update_token_counters(self, nt_in, nt_out): self.total_tokens_in += nt_in self.total_tokens_out += nt_out def update_final_response(self, response_candidate, WFO_record, GEO_record, paths, path_to_crop, nt_in, nt_out): filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper = paths # Saving the JSON and XLSX files with the response and updating the final JSON response if response_candidate is not None: final_JSON_response_updated = self.save_json_and_xlsx(response_candidate, WFO_record, GEO_record, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out) return final_JSON_response_updated, WFO_record, GEO_record else: final_JSON_response_updated = self.save_json_and_xlsx(response_candidate, WFO_record, GEO_record, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out) return final_JSON_response_updated, WFO_record, GEO_record def log_completion_info(self, final_JSON_response): self.logger.info(f'Formatted JSON\n{final_JSON_response}') self.logger.info(f'Finished API calls\n') def update_progress_report_final(self, progress_report): if progress_report is not None: progress_report.reset_batch("Batch Complete") def parse_final_json_response(self, final_JSON_response): try: return json.loads(final_JSON_response.strip('```').replace('json\n', '', 1).replace('json', '', 1)) except: return final_JSON_response def generate_paths(self, path_to_crop, i): filename_without_extension = os.path.splitext(os.path.basename(path_to_crop))[0] txt_file_path = os.path.join(self.Dirs.transcription_ind, filename_without_extension + '.json') txt_file_path_OCR = os.path.join(self.Dirs.transcription_ind_OCR, filename_without_extension + '.json') txt_file_path_OCR_bounds = os.path.join(self.Dirs.transcription_ind_OCR_bounds, filename_without_extension + '.json') jpg_file_path_OCR_helper = os.path.join(self.Dirs.transcription_ind_OCR_helper, filename_without_extension + '.jpg') self.logger.info(f'Working on {i+1}/{len(self.img_paths)} --- {filename_without_extension}') return filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper def save_json_and_xlsx(self, response, WFO_record, GEO_record, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out): if response is None: response = self.JSON_dict_structure # Insert 'filename' as the first key response = {'filename': filename_without_extension, **{k: v for k, v in response.items() if k != 'filename'}} self.write_json_to_file(txt_file_path, response) # Then add the null info to the spreadsheet response_null = self.create_null_row(filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper) self.add_data_to_excel_from_response(self.path_transcription, response_null, WFO_record, GEO_record, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in=0, nt_out=0) ### Set completed JSON else: response = self.clean_catalog_number(response, filename_without_extension) self.write_json_to_file(txt_file_path, response) # add to the xlsx file self.add_data_to_excel_from_response(self.path_transcription, response, WFO_record, GEO_record, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out) return response def process_specimen_batch(self, progress_report, json_report, is_real_run=False): if not self.has_key: self.logger.error(f'No API key found for {self.version_name}') raise Exception(f"No API key found for {self.version_name}") try: if is_real_run: progress_report.update_overall(f"Transcribing Labels") final_json_response, final_WFO_record, final_GEO_record, total_tokens_in, total_tokens_out = self.send_to_LLM(self.is_azure, progress_report, json_report, self.model_name) return final_json_response, final_WFO_record, final_GEO_record, total_tokens_in, total_tokens_out except Exception as e: self.logger.error(f"LLM call failed in process_specimen_batch: {e}") if progress_report is not None: progress_report.reset_batch(f"Batch Failed") self.close_logger_handlers() raise def close_logger_handlers(self): for handler in self.logger.handlers[:]: handler.close() self.logger.removeHandler(handler) def process_specimen_batch_OCR_test(self, path_to_crop): for img_filename in os.listdir(path_to_crop): img_path = os.path.join(path_to_crop, img_filename) self.OCR, self.bounds, self.text_to_box_mapping = detect_text(img_path) def space_saver(cfg, Dirs, logger): dir_out = cfg['leafmachine']['project']['dir_output'] run_name = Dirs.run_name path_project = os.path.join(dir_out, run_name) if cfg['leafmachine']['project']['delete_temps_keep_VVE']: logger.name = '[DELETE TEMP FILES]' logger.info("Deleting temporary files. Keeping files required for VoucherVisionEditor.") delete_dirs = ['Archival_Components', 'Config_File'] for d in delete_dirs: path_delete = os.path.join(path_project, d) if os.path.exists(path_delete): shutil.rmtree(path_delete) elif cfg['leafmachine']['project']['delete_all_temps']: logger.name = '[DELETE TEMP FILES]' logger.info("Deleting ALL temporary files!") delete_dirs = ['Archival_Components', 'Config_File', 'Original_Images', 'Cropped_Images'] for d in delete_dirs: path_delete = os.path.join(path_project, d) if os.path.exists(path_delete): shutil.rmtree(path_delete) # Delete the transctiption folder, but keep the xlsx transcription_path = os.path.join(path_project, 'Transcription') if os.path.exists(transcription_path): for item in os.listdir(transcription_path): item_path = os.path.join(transcription_path, item) if os.path.isdir(item_path): # if the item is a directory if os.path.exists(item_path): shutil.rmtree(item_path) # delete the directory