fixed entity merging issue

handwritting_detection.py

@@ -36,6 +36,6 @@ def DetectHandwritting(image):
     cpy = image.copy()
     handwritten_parts = []
     for prediction in result['predictions']:
-        cpy = draw_rectangle(cpy, **prediction)
         handwritten_parts.append(crop_image(cpy, **prediction))
+        cpy = draw_rectangle(cpy, **prediction)
     return cpy, handwritten_parts

main.py

@@ -49,7 +49,8 @@ async def ProcessDocument(file: UploadFile):
   
   try:
    reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
-  except:
+  except Exception as e:
+    print(e)
     raise HTTPException(status_code=400, detail="Relation extraction failed")
   return reOutput
 
@@ -66,9 +67,14 @@ async def ProcessDocument(file: str = Form(...)):
     raise HTTPException(status_code=400, detail="Cannot apply OCR to the image")
   try:
     tokenClassificationOutput, img_size = LabelTokens(ocr_df, image)
+  except Exception as e:
+    print(e)
+    raise HTTPException(status_code=400, detail="Entity identification failed")
+  try:
     reOutput = ExtractRelations(tokenClassificationOutput, ocr_df, img_size)
-  except:
-    raise HTTPException(status_code=400, detail="Invalid Image")
+  except Exception as e:
+    print(e)
+    raise HTTPException(status_code=400, detail="Relation extraction failed")
   return reOutput
 
 def ApplyOCR(content):
@@ -112,7 +118,6 @@ def LabelTokens(ocr_df, image):
   return {"token_labels": token_labels, "input_ids": input_ids, "bbox":bbox, "attention_mask":attention_mask}, image.size
 
 def ExtractRelations(tokenClassificationOutput, ocr_df, img_size):
-  print(tokenClassificationOutput)
   token_labels = tokenClassificationOutput['token_labels']
   input_ids = tokenClassificationOutput['input_ids']
   attention_mask = tokenClassificationOutput["attention_mask"]

token_classification.py

@@ -195,10 +195,9 @@ def createEntities(model, predictions, input_ids, ocr_df, tokenizer, img_size, b
     for i in range(len(df_words)):
         if i not in visited:
             merged_taggings.append(dfs(i,[], width, height, visited, df_words))
-
     merged_words = []
     for i,merged_tagging in enumerate(merged_taggings):
-        if ((len(merged_tagging) > 1) or (merged_tagging['label']=='ANSWER')):
+        if ((len(merged_tagging) > 1)) or (merged_tagging[0]['label'] == 'ANSWER'):
             new_word = {}
             merging_word = " ".join([word['text'] for word in merged_tagging])
             merging_box = [merged_tagging[0]['box'][0]-5,merged_tagging[0]['box'][1]-10,merged_tagging[-1]['box'][2]+5,merged_tagging[-1]['box'][3]+10]