Spaces:
Runtime error
Runtime error
| # -*- 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("alitavanaali/music_layoutlmv3_model") | |
| feature_extractor = LayoutLMv3FeatureExtractor(apply_ocr=True,lang=lang) | |
| 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): | |
| 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='eng' | |
| width, height = image.size | |
| 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 = [] # related to more than 512 tokens | |
| 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() | |
| #{0: 'document number', 1: 'elemento pn', 2: 'nombre del responsabile', 3: 'fecha', 4: 'internal reference', 5: 'others'} | |
| #id2label = {0: 'song name', 1: 'artist', 2: 'year', 3: 'album', 4: 'genres', 5: 'song writer', 6: 'lyrics', 7: 'others'} | |
| 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 = [] | |
| # draw bboxes on image | |
| for ix, (prediction, box) in enumerate(zip(final_preds, final_bbox)): | |
| predicted_label = iob_to_label(prediction).lower() | |
| if (predicted_label != 'others'): | |
| 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: | |
| flag = False | |
| for b, p, w in zip(bb, pr, wr): | |
| if (p == 'others'): | |
| #print('others') | |
| #print('-------') | |
| continue | |
| elif (box==b): # we shouldn't check each item with itself | |
| #print('itself') | |
| #print('--------') | |
| continue | |
| else: | |
| XMIN, YMIN, XMAX, YMAX = box | |
| xmin, ymin, xmax, ymax = b | |
| #print('word: {} , w:{}'.format(word, w)) | |
| intsc = intersect([XMIN, YMIN, XMAX+threshold, YMAX], [xmin-threshold, ymin, xmax, ymax]) | |
| if (intsc != 0 and pred==p): | |
| flag = True | |
| #print('there is intersect') | |
| # if(abs(XMAX - xmin) < treshold and abs(YMIN - ymin) < 10): | |
| # we have to check if there is any intersection between box and all the values in final_bbox list | |
| # because if we have updated it before, now we have to update in final_bbox | |
| #print(final_bbox) | |
| print(*final_bbox, sep=",") | |
| merged_box = [ | |
| min(XMIN, xmin), | |
| min(YMIN, ymin), | |
| max(XMAX, xmax), | |
| max(YMAX, ymax) | |
| ] | |
| merged_words = word + ' ' + w | |
| # add to final_bbox | |
| wasAvailable = False | |
| for id, fbox in enumerate(final_bbox): | |
| if (intersect(box, fbox) != 0 and pred==final_preds[id]): | |
| #print('added before!') | |
| # box is inside another processed box, so we have to update it | |
| wasAvailable = True | |
| merged_box = [ | |
| min(fbox[0], min(XMIN, xmin)), | |
| min(fbox[1], min(YMIN, ymin)), | |
| max(fbox[2], max(XMAX, xmax)), | |
| max(fbox[3], max(YMAX, ymax)) | |
| ] | |
| final_bbox[id] = merged_box | |
| final_words[id] = final_words[id] + ' ' + w | |
| break | |
| if (not wasAvailable): | |
| # there was no intersect, bbox is not added before | |
| #print('not added before, so we add merged box!') | |
| final_bbox.append(merged_box) | |
| final_preds.append(pred) | |
| final_words.append(merged_words) | |
| '''else: | |
| print() | |
| final_bbox.append(box) | |
| final_preds.append(pred) | |
| final_words.append(word)''' | |
| if (flag == False): | |
| #print('flag is false, word: {} added'.format(word)) | |
| # there is no intersect between word and the others | |
| # we will check for last time if box is inside the others, because if the word is last word (like Juan + Mulian + Alexander) (Alexander) | |
| # it is added before but it has not intersection with others, so we will check to prevent | |
| for id, fbox in enumerate(final_bbox): | |
| if (intersect(box, fbox) != 0 and pred==final_preds[id]): | |
| flag = True | |
| if (not flag): | |
| final_bbox.append(box) | |
| final_preds.append(pred) | |
| final_words.append(word) | |
| return final_bbox, final_preds, final_words | |
| def createDataframe(preds, words): | |
| df = pd.DataFrame(columns = ['song name', 'artist', 'year', 'album', 'genres', 'song writer', 'lyrics', 'others']) | |
| if (len(preds) > 0): | |
| 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, 30) | |
| _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) | |
| print(preds) | |
| print(words) | |
| 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) |