BilalHasan commited on
Commit
1fd5d04
1 Parent(s): 6cf0b65

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +108 -0
app.py ADDED
@@ -0,0 +1,108 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ image = cv.imread("/content/Demo2.png")
2
+ def preprocess_image(image):
3
+ gray_image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
4
+ ret, bin_image = cv.threshold(gray_image, 127, 255, cv.THRESH_OTSU)
5
+
6
+ bin_image = cv.copyMakeBorder(bin_image, int(0.10 * image.shape[0]), int(0.05 * image.shape[0]), int(0.05 * image.shape[1]), int(0.10 * image.shape[1]), cv.BORDER_CONSTANT, value=(255, 255, 255))
7
+ return bin_image
8
+
9
+ bin_image = preprocess_image(image)
10
+
11
+ def split_image_into_lines(image):
12
+ lines = []
13
+ while (image.shape[0] > 20):
14
+ flag1 = 0
15
+ flag2 = 0
16
+
17
+ for i in range(image.shape[0]):
18
+ if flag1 == 0:
19
+ for j in range(image.shape[1]):
20
+ pixel_value = image[i][j]
21
+ if (pixel_value == 0) & (flag1 == 0):
22
+ start = i
23
+ flag1 = 1
24
+ flag2 = 1
25
+ if flag2 == 1:
26
+ num_white_pixels = np.sum(image[i + 1] == 255)
27
+ if (num_white_pixels > 0.98 * image.shape[1]):
28
+ end = i + 1
29
+ break
30
+
31
+
32
+ line = image[int(start - 0.2 * (end - start + 1)): int(end + 1 + 0.2 * (end - start + 1))][:]
33
+ if line.shape[0] > 20:
34
+ line_rgb = cv.cvtColor(line, cv.COLOR_GRAY2RGB)
35
+ lines.append(line_rgb)
36
+
37
+ pads = 255 * np.ones((20, image.shape[1]), dtype='uint8')
38
+ new_image = image[int(end + 2 -(0.2 * (end - start + 1))):][:]
39
+ new_image = np.concatenate((pads, new_image))
40
+ image = new_image
41
+
42
+
43
+ return lines
44
+
45
+ lines = split_image_into_lines(bin_image)
46
+
47
+
48
+ def generate_text(line):
49
+ pixel_values = processor(images=line, return_tensors="pt").pixel_values
50
+ generated_ids = model.generate(pixel_values, max_new_tokens=50)
51
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
52
+ return generated_text
53
+
54
+ with ProcessPoolExecutor() as executor:
55
+ results = ' '.join(executor.map(generate_text, lines))
56
+ #improve results with llm
57
+
58
+ client = OpenAI()
59
+
60
+ completion = client.chat.completions.create(
61
+ model="gpt-4o",
62
+ messages=[
63
+ {
64
+ "role": "user",
65
+
66
+ "content": f"I have a string that was extracted from an image of handwritten text. The extraction process introduced minor grammatical, spelling, and punctuation errors. Please carefully review the text below and make any necessary corrections to improve readability and accuracy while preserving the original meaning. Do not change the content or style beyond necessary corrections. Return the corrected text only without adding any headings, explanations, or extra formatting. Text: {results}"
67
+ }
68
+ ]
69
+ )
70
+
71
+ improved_text = completion.choices[0].message.content
72
+
73
+
74
+
75
+ def put_text(text, font, font_scale, color, thickness, max_width, out_image_width, top_margin):
76
+ words = text.split(" ")
77
+ lines = []
78
+ current_line = ""
79
+
80
+ for word in words:
81
+ if cv.getTextSize(current_line + " " + word, font, font_scale, thickness)[0][0] <= (max_width * out_image_width):
82
+ current_line += " " + word
83
+ else:
84
+ lines.append(current_line)
85
+ current_line = word
86
+
87
+ lines.append(current_line)
88
+
89
+ out_image_height = sum([cv.getTextSize(line, font, font_scale, thickness)[0][1] for line in lines]) + 2 * top_margin + 20 * (len(lines) - 1) #20 is the gap between two consecutive lines
90
+
91
+ out_image = 255 * (np.ones((out_image_height, out_image_width, 3), dtype=np.uint8))
92
+
93
+ top = top_margin
94
+ for line in lines:
95
+ cv.putText(out_image, line.strip(), (int(((1 - max_width) * out_image_width) / 2), top), font, font_scale, 0, thickness, lineType=cv.LINE_AA)
96
+ top += cv.getTextSize(line.strip(), font, font_scale, thickness)[0][1] + 20
97
+
98
+ return out_image
99
+
100
+ font = cv.FONT_HERSHEY_DUPLEX
101
+ font_scale = 2
102
+ color = 0
103
+ thickness = 2
104
+ max_width = 0.9
105
+ out_image_width = 1500
106
+ top_margin = 100
107
+
108
+ out_image = put_text(improved_text, font, font_scale, color, thickness, max_width, out_image_width, top_margin)