# -*- coding: utf-8 -*- """Untitled1.ipynb Automatically generated by Colaboratory. Original file is located at https://colab.research.google.com/drive/1J4fCr7TGzdFvkCeikMAQ5af5ml2Q83W0 """ import os os.system('pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu') import os, glob, fitz import cv2 import os import PIL import torch import pandas as pd import numpy as np import pandas as pd import gradio as gr from tqdm import tqdm from PIL import Image as im from scipy import ndimage from difflib import SequenceMatcher from itertools import groupby from datasets import load_metric from datasets import load_dataset from datasets.features import ClassLabel from transformers import AutoProcessor from PIL import Image, ImageDraw, ImageFont from transformers import AutoModelForTokenClassification from transformers.data.data_collator import default_data_collator from datasets import Features, Sequence, ClassLabel, Value, Array2D, Array3D from transformers import LayoutLMv3ForTokenClassification,LayoutLMv3FeatureExtractor # define id2label id2label = {0: 'song name', 1: 'artist', 2: 'year', 3: 'album', 4: 'genres', 5: 'song writer', 6: 'lyrics', 7: 'others'} custom_config = r'--oem 3 --psm 6' # lang='eng+deu+ita+chi_sim' lang='spa' label_ints = np.random.randint(0,len(PIL.ImageColor.colormap.items()),42) label_color_pil = [k for k,_ in PIL.ImageColor.colormap.items()] label_color = [label_color_pil[i] for i in label_ints] label2color = {} for k,v in id2label.items(): if v[:2] == '': label2color['o']=label_color[k] else: label2color[v[2:]]=label_color[k] processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base", apply_ocr=True,lang=lang) model = AutoModelForTokenClassification.from_pretrained("atatavana/layoutlm_manifesto_bigdataset") def unnormalize_box(bbox, width, height): #print('shape is: ', np.asarray(bbox).shape, ' and box has values: ', bbox) return [ width * (bbox[0] / 1000), height * (bbox[1] / 1000), width * (bbox[2] / 1000), height * (bbox[3] / 1000), ] def iob_to_label(label): if label == 0: return 'song name' if label == 1: return 'artist' if label == 2: return 'year' if label == 3: return 'album' if label == 4: return 'genres' if label == 5: return 'song writer' if label == 6: return 'lyrics' if label == 7: return 'others' # this method will detect if there is any intersect between two boxes or not def intersect(w, z): x1 = max(w[0], z[0]) #190 | 881 | 10 y1 = max(w[1], z[1]) #90 | 49 | 273 x2 = min(w[2], z[2]) #406 | 406 | 1310 y2 = min(w[3], z[3]) #149 | 703 | 149 if (x1 > x2 or y1 > y2): return 0 else: # because sometimes in annotating, it is possible to overlap rows or columns by mistake # for very small pixels, we check a threshold to delete them area = (x2-x1) * (y2-y1) if (area > 0): #500 is minumum accepted area return [int(x1), int(y1), int(x2), int(y2)] else: return 0 def process_image(image): custom_config = r'--oem 3 --psm 6' # lang='eng+deu+ita+chi_sim' lang='spa' width, height = image.size feature_extractor = LayoutLMv3FeatureExtractor(apply_ocr=True,lang=lang) encoding_feature_extractor = feature_extractor(image, return_tensors="pt",truncation=True) words, boxes = encoding_feature_extractor.words, encoding_feature_extractor.boxes custom_config = r'--oem 3 --psm 6' # encode inference_image = [image.convert("RGB")] encoding = processor(inference_image , truncation=True, return_offsets_mapping=True, return_tensors="pt", padding="max_length", stride =128, max_length=512, return_overflowing_tokens=True) offset_mapping = encoding.pop('offset_mapping') overflow_to_sample_mapping = encoding.pop('overflow_to_sample_mapping') # change the shape of pixel values x = [] for i in range(0, len(encoding['pixel_values'])): x.append(encoding['pixel_values'][i]) x = torch.stack(x) encoding['pixel_values'] = x # forward pass outputs = model(**encoding) # get predictions predictions = outputs.logits.argmax(-1).squeeze().tolist() token_boxes = encoding.bbox.squeeze().tolist() # only keep non-subword predictions preds = [] l_words = [] bboxes = [] token_section_num = [] if (len(token_boxes) == 512): predictions = [predictions] token_boxes = [token_boxes] for i in range(0, len(token_boxes)): for j in range(0, len(token_boxes[i])): #print(np.asarray(token_boxes[i][j]).shape) unnormal_box = unnormalize_box(token_boxes[i][j], width, height) #print('prediction: {} - box: {} - word:{}'.format(predictions[i][j], unnormal_box, processor.tokenizer.decode(encoding["input_ids"][i][j]))) if (np.asarray(token_boxes[i][j]).shape != (4,)): continue elif (token_boxes[i][j] == [0, 0, 0, 0] or token_boxes[i][j] == 0): #print('zero found!') continue # if bbox is available in the list, just we need to update text elif (unnormal_box not in bboxes): preds.append(predictions[i][j]) l_words.append(processor.tokenizer.decode(encoding["input_ids"][i][j])) bboxes.append(unnormal_box) token_section_num.append(i) else: # we have to update the word _index = bboxes.index(unnormal_box) if (token_section_num[_index] == i): # check if they're in a same section or not (documents with more than 512 tokens will divide to seperate # parts, so it's possible to have a word in both of the pages and we have to control that repetetive words # HERE: because they're in a same section, so we can merge them safely l_words[_index] = l_words[_index] + processor.tokenizer.decode(encoding["input_ids"][i][j]) else: continue return bboxes, preds, l_words, image def visualize_image(final_bbox, final_preds, l_words, image): draw = ImageDraw.Draw(image) font = ImageFont.load_default() label2color = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue', 'lyrics':'purple', 'others': 'white'} l2l = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue','lyrics':'purple', 'others':'white'} f_labels = {'song name':'red', 'artist':'blue', 'year':'black', 'album':'green', 'genres':'brown', 'song writer':'blue','lyrics':'purple', 'others':'white'} json_df = [] for ix, (prediction, box) in enumerate(zip(final_preds, final_bbox)): predicted_label = iob_to_label(prediction).lower() draw.rectangle(box, outline=label2color[predicted_label]) draw.text((box[0]+10, box[1]-10), text=predicted_label, fill=label2color[predicted_label], font=font) json_dict = {} json_dict['TEXT'] = l_words[ix] json_dict['LABEL'] = f_labels[predicted_label] json_df.append(json_dict) return image, json_df def mergeCloseBoxes(pr, bb, wr, threshold): idx = 0 final_bbox =[] final_preds =[] final_words=[] for box, pred, word in zip(bb, pr, wr): if (pred=='others'): continue else: final_bbox.append(box) final_preds.append(pred) final_words.append(word) for b, p, w in zip(bb, pr, wr): if (p == 'others'): continue elif (box==b): # we shouldn't check each item with itself continue else: XMIN, YMIN, XMAX, YMAX = box xmin, ymin, xmax, ymax = b intsc = intersect([XMIN, YMIN, XMAX+threshold, YMAX], [xmin-threshold, ymin, xmax, ymax]) if (intsc != 0 and pred==p): #if(abs(XMAX - xmin) < treshold and abs(YMIN - ymin) < 10): if(box in final_bbox): final_bbox[idx]= [XMIN, min(YMIN, ymin), xmax, max(YMAX, ymax)] final_words[idx] = word + ' ' + w continue print('box: {}, label: {} is close to b:{} with this p:{}--> {}'.format(box, pred, b, p, word + ' ' + w)) idx = idx +1 return final_bbox, final_preds, final_words def createDataframe(preds, words): df = pd.DataFrame(columns = ['container id' ,'seal number', 'container quantity', 'container type', 'package quantity', 'tare', 'weight']) flag_label = preds[0] #print(preds) #print(words) #print('@@@@@') #print(flag_label) row_number = -1 for i in range(len(preds)): #print('i is: {}'.format(i)) if (preds[i] == flag_label): row_number = row_number + 1 df.at[row_number, preds[i]] = words[i] #print('row number is: {}'.format(row_number)) continue else: #print('row_number {} is <= of df.shape {}'.format(row_number, df.shape[0])) #print(pd.isna(df[preds[i]].iloc[row_number])) #print(pd.isna(df[preds[i]].iloc[row_number])) if(pd.isna(df[preds[i]].iloc[row_number])): df.at[row_number, preds[i]] = words[i] else: row_number = row_number + 1 df.at[row_number, preds[i]] = words[i] return df def isInside(w, z): # return True if w is inside z, if z is inside w return false if(w[0] >= z[0] and w[1] >= z[1] and w[2] <= z[2] and w[3] <= z[3]): return True return False def removeSimilarItems(final_bbox, final_preds, final_words): _bb =[] _pp=[] _ww=[] for i in range(len(final_bbox)): _bb.append(final_bbox[i]) _pp.append(final_preds[i]) _ww.append(final_words[i]) for j in range(len(final_bbox)): if (final_bbox[i] == final_bbox[j]): continue elif (isInside(final_bbox[i], final_bbox[j]) and final_preds[i]==final_preds[j] ): # box i is inside box j, so we have to remove it #print('box[i]: {} is inside box[j]:{}'.format(final_bbox[i], final_bbox[j])) _bb = _bb[:-1] _pp = _pp[:-1] _ww = _ww[:-1] continue return _bb, _pp, _ww #[45.604, 2309.811, 66.652, 2391.6839999999997] def process_form(preds, words, bboxes): final_bbox, final_preds, final_words = mergeCloseBoxes(preds, bboxes, words, 70) _bbox, _preds, _words = removeSimilarItems(final_bbox, final_preds, final_words) # convert float list to int _bbox = [[int(x) for x in item ] for item in _bbox] # creat data object for sorting data = [] for index in range(len(_bbox)): data.append((_bbox[index], _preds[index], _words[index])) # sorting by the height of the page sorted_list = sorted( data, key=lambda x: x[0][1] ) _bbox = [item[0] for item in sorted_list] _preds = [item[1] for item in sorted_list] _words = [item[2] for item in sorted_list] return _bbox, _preds, _words def mergeImageVertical(a): list_im = a imgs = [ Image.open(i) for i in list_im ] # pick the image which is the smallest, and resize the others to match it (can be arbitrary image shape here) min_shape = sorted( [(np.sum(i.size), i.size ) for i in imgs])[0][1] imgs_comb = np.hstack([i.resize(min_shape) for i in imgs]) # for a vertical stacking it is simple: use vstack imgs_comb = np.vstack([i.resize(min_shape) for i in imgs]) imgs_comb = Image.fromarray( imgs_comb) imgs_comb.save( 'Trifecta_vertical.jpg' ) return imgs_comb def completepreprocess(pdffile): myDataFrame = pd.DataFrame() a=[] doc = fitz.open(pdffile) for i in range(0,len(doc)): page = doc.load_page(i) zoom = 2 # zoom factor mat = fitz.Matrix(zoom, zoom) pix = page.get_pixmap(matrix = mat,dpi = 200) t=pix.save("page"+str(i)+".jpg") images = Image.open("page"+str(i)+".jpg") image = images.convert("RGB") bbox, preds, words, image = process_image(image) im, df = visualize_image(bbox, preds, words, image) im1 = im.save("page"+str(i)+".jpg") a.append("page"+str(i)+".jpg") pred_list = [] for number in preds: pred_list.append(iob_to_label(number)) _bbox, _preds, _words = process_form(pred_list, words, bbox) print('page: ' + str(i) + ' ' + str(len(_preds))+ ' ' + str(len(_words))) df = createDataframe(_preds, _words) myDataFrame=myDataFrame.append(df) im2=mergeImageVertical(a) return im2,myDataFrame title = "Interactive demo: Music Information Extraction model" description = "Music Information Extraction - We used Microsoft’s LayoutLMv3 trained on Our Music Dataset through csv's to predict the labels. To use it, simply upload a PDF or use the example PDF below and click ‘Submit’. Results will show up in a few seconds. If you want to make the output bigger, right-click on it and select ‘Open image in new tab’.Train =16 ,Test =7" css = """.output_image, .input_image {height: 600px !important}""" #examples = [["461BHH69.PDF"],["AP-481-RF.PDF"],["DP-095-ML.PDF"],["DQ-231-LL.PDF"],["FK-941-ET.PDF"], ["FL-078-NH.PDF"] # ,["14ZZ69.PDF"],["74BCA69.PDF"],["254BEG69.PDF"],["761BJQ69.PDF"],["AB-486-EH.PDF"],["AZ-211-ZA.PDF"], ["CY-073-YV.PDF"]] # ["744BJQ69.PDF"], ['tarros_2.jpg'], examples = [['test1.jpg'], ['doc1.pdf'], ['doc1.2.pdf']] iface = gr.Interface(fn=completepreprocess, #inputs=gr.inputs.Image(type="pil",optional=True,label="upload file"), inputs=gr.File(label="PDF"), #inputs=gr.inputs.Image(type="pil") outputs=[gr.outputs.Image(type="pil", label="annotated image"),"dataframe"] , title=title, description=description, examples=examples, css=css, analytics_enabled = True, enable_queue=True) iface.launch(inline=False , debug=True)