AuditEdge commited on
Commit
f8afc9b
·
1 Parent(s): 3cbb02d

doc upload option added

Browse files
Files changed (2) hide show
  1. app.py +40 -6
  2. utils.py +110 -0
app.py CHANGED
@@ -4,13 +4,15 @@ from typing import Dict
4
  import os
5
  import shutil
6
  import logging
7
-
8
  import torch
9
  from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
10
 
11
  from dotenv import load_dotenv
12
  import os
13
 
 
 
14
  # Load .env file
15
  load_dotenv()
16
 
@@ -33,6 +35,8 @@ aadhar_model = LayoutLMv3ForTokenClassification.from_pretrained(
33
  aadhar_model,
34
  use_auth_token=HUGGINGFACE_AUTH_TOKEN
35
  )
 
 
36
  aadhar_model = aadhar_model.to(device)
37
 
38
  # pan model
@@ -40,6 +44,8 @@ pan_model = "AuditEdge/doc_ocr_p" # Replace with your fine-tuned model if appli
40
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
41
  print(f"Using device: {device}")
42
 
 
 
43
  # Load the processor (tokenizer + image processor)
44
  processor_pan = LayoutLMv3Processor.from_pretrained(
45
  pan_model,
@@ -113,7 +119,9 @@ app.add_middleware(
113
 
114
  # Configure directories
115
  UPLOAD_FOLDER = './uploads/'
 
116
  os.makedirs(UPLOAD_FOLDER, exist_ok=True) # Ensure the main upload folder exists
 
117
 
118
  UPLOAD_DIRS = {
119
  "aadhar_file": "uploads/aadhar/",
@@ -122,9 +130,22 @@ UPLOAD_DIRS = {
122
  "gst_file": "uploads/gst/",
123
  }
124
 
 
 
 
 
 
 
 
 
125
  # Ensure individual directories exist
126
  for dir_path in UPLOAD_DIRS.values():
127
  os.makedirs(dir_path, exist_ok=True)
 
 
 
 
 
128
 
129
  # Logger configuration
130
  logging.basicConfig(level=logging.INFO)
@@ -217,14 +238,27 @@ async def aadhar_ocr(
217
  # Log received files
218
  logging.info(f"Received files: {list(file_paths.keys())}")
219
  print("file_paths",file_paths)
220
- import sys
221
- # sys.exit()
222
-
 
 
 
 
 
 
 
 
 
 
223
  # Perform inference
224
- result = perform_inference(file_paths)
225
 
226
  return {"status": "success", "result": result}
227
 
228
  except Exception as e:
229
  logging.error(f"Error processing files: {e}")
230
- raise HTTPException(status_code=500, detail="Internal Server Error")
 
 
 
 
4
  import os
5
  import shutil
6
  import logging
7
+
8
  import torch
9
  from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
10
 
11
  from dotenv import load_dotenv
12
  import os
13
 
14
+ from utils import doc_processing
15
+
16
  # Load .env file
17
  load_dotenv()
18
 
 
35
  aadhar_model,
36
  use_auth_token=HUGGINGFACE_AUTH_TOKEN
37
  )
38
+
39
+
40
  aadhar_model = aadhar_model.to(device)
41
 
42
  # pan model
 
44
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
45
  print(f"Using device: {device}")
46
 
47
+
48
+
49
  # Load the processor (tokenizer + image processor)
50
  processor_pan = LayoutLMv3Processor.from_pretrained(
51
  pan_model,
 
119
 
120
  # Configure directories
121
  UPLOAD_FOLDER = './uploads/'
122
+ processing_folder = "./processed_images"
123
  os.makedirs(UPLOAD_FOLDER, exist_ok=True) # Ensure the main upload folder exists
124
+ os.makedirs(processing_folder,exist_ok=True)
125
 
126
  UPLOAD_DIRS = {
127
  "aadhar_file": "uploads/aadhar/",
 
130
  "gst_file": "uploads/gst/",
131
  }
132
 
133
+ process_dirs = {
134
+ "aadhar_file": "processed_images/aadhar/",
135
+ "pan_file": "processed_images/pan/",
136
+ "cheque_file": "processed_images/cheque/",
137
+ "gst_file": "processed_images/gst/",
138
+
139
+ }
140
+
141
  # Ensure individual directories exist
142
  for dir_path in UPLOAD_DIRS.values():
143
  os.makedirs(dir_path, exist_ok=True)
144
+
145
+ for dir_path in process_dirs.values():
146
+ os.makedirs(dir_path, exist_ok=True)
147
+
148
+
149
 
150
  # Logger configuration
151
  logging.basicConfig(level=logging.INFO)
 
238
  # Log received files
239
  logging.info(f"Received files: {list(file_paths.keys())}")
240
  print("file_paths",file_paths)
241
+
242
+ files = {}
243
+ for key, value in file_paths.items():
244
+ name = value.split("/")[-1].split(".")[0]
245
+ id_type = key.split("_")[0]
246
+ doc_type = value.split("/")[-1].split(".")[1]
247
+ f_path = value
248
+ preprocessing = doc_processing(name,id_type,doc_type,f_path)
249
+ response = preprocessing.process()
250
+ files[key] = response["output_p"]
251
+ print("response",response)
252
+
253
+
254
  # Perform inference
255
+ result = perform_inference(files)
256
 
257
  return {"status": "success", "result": result}
258
 
259
  except Exception as e:
260
  logging.error(f"Error processing files: {e}")
261
+ # raise HTTPException(status_code=500, detail="Internal Server Error")
262
+ return {"status":400}
263
+
264
+
utils.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import fitz
2
+ from PIL import Image
3
+
4
+ class doc_processing:
5
+
6
+ def __init__(self, name, id_type, doc_type, f_path):
7
+
8
+ self.name = name
9
+ self.id_type = id_type
10
+ self.doc_type = doc_type
11
+ self.f_path = f_path
12
+ # self.o_path = o_path
13
+
14
+
15
+ def pdf_to_image_scale(self):
16
+ pdf_document = fitz.open(self.f_path)
17
+ if self.id_type == "gst":
18
+ page_num = 2
19
+ else:
20
+ page_num = 0
21
+
22
+ page = pdf_document.load_page(page_num)
23
+ pix = page.get_pixmap() # Render page as a pixmap (image)
24
+
25
+ # Convert pixmap to PIL Image
26
+ image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
27
+
28
+ original_width, original_height = image.size
29
+
30
+ print("original_width",original_width)
31
+ print("original_height",original_height)
32
+
33
+
34
+ new_width = (1000 / original_width) * original_width
35
+ new_height = (1000 / original_height) * original_height
36
+
37
+ print("new_width",new_width)
38
+ print("new_height",new_height)
39
+ # new_width =
40
+ # new_height =
41
+ image.resize((int(new_width), int(new_height)), Image.Resampling.LANCZOS)
42
+ output_path = "processed_images/{}/{}.jpeg".format(self.id_type,self.name)
43
+ image.save(output_path)
44
+ return {"success":200,"output_p":output_path}
45
+
46
+
47
+ def scale_img(self):
48
+
49
+ image = Image.open(self.f_path).convert("RGB")
50
+ original_width, original_height = image.size
51
+
52
+ print("original_width",original_width)
53
+ print("original_height",original_height)
54
+
55
+
56
+ new_width = (1000 / original_width) * original_width
57
+ new_height = (1000 / original_height) * original_height
58
+
59
+ print("new_width",new_width)
60
+ print("new_height",new_height)
61
+ # new_width =
62
+ # new_height =
63
+ image.resize((int(new_width), int(new_height)), Image.Resampling.LANCZOS)
64
+ output_path = "processed_images/{}/{}.jpeg".format(self.id_type,self.name)
65
+ image.save(output_path)
66
+ return {"success":200,"output_p":output_path}
67
+
68
+ def process(self):
69
+ if self.doc_type == "pdf":
70
+ response = self.pdf_to_image_scale()
71
+ else:
72
+ response = self.scale_img()
73
+
74
+ return response
75
+
76
+
77
+
78
+
79
+
80
+ # files = {
81
+ # "aadhar_file": "/home/javmulla/model_one/test_images_aadhar/test_two.jpg",
82
+ # "pan_file": "/home/javmulla/model_one/test_images_pan/6ea33087.jpeg",
83
+ # "cheque_file": "/home/javmulla/model_one/test_images_cheque/0f81678a.jpeg",
84
+ # "gst_file": "/home/javmulla/model_one/test_images_gst/0a52fbcb_page3_image_0.jpg"
85
+ # }
86
+
87
+
88
+ # files = {
89
+ # "aadhar_file": "/home/javmulla/model_one/test_images_aadhar/test_two.jpg",
90
+ # "pan_file": "/home/javmulla/model_one/test_images_pan/6ea33087.jpeg",
91
+ # "cheque_file": "/home/javmulla/model_one/test_images_cheque/0f81678a.jpeg",
92
+ # "gst_file": "test_Images_folder/gst/e.pdf"
93
+ # }
94
+
95
+ # for key, value in files.items():
96
+ # name = value.split("/")[-1].split(".")[0]
97
+ # id_type = key.split("_")[0]
98
+ # doc_type = value.split("/")[-1].split(".")[1]
99
+ # f_path = value
100
+ # preprocessing = doc_processing(name,id_type,doc_type,f_path)
101
+ # response = preprocessing.process()
102
+ # print("response",response)
103
+
104
+
105
+
106
+
107
+
108
+ # id_type, doc_type, f_path
109
+
110
+