Spaces:
Running
Running
| image = cv.imread("/content/Demo2.png") | |
| def preprocess_image(image): | |
| gray_image = cv.cvtColor(image, cv.COLOR_BGR2GRAY) | |
| ret, bin_image = cv.threshold(gray_image, 127, 255, cv.THRESH_OTSU) | |
| 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)) | |
| return bin_image | |
| bin_image = preprocess_image(image) | |
| def split_image_into_lines(image): | |
| lines = [] | |
| while (image.shape[0] > 20): | |
| flag1 = 0 | |
| flag2 = 0 | |
| for i in range(image.shape[0]): | |
| if flag1 == 0: | |
| for j in range(image.shape[1]): | |
| pixel_value = image[i][j] | |
| if (pixel_value == 0) & (flag1 == 0): | |
| start = i | |
| flag1 = 1 | |
| flag2 = 1 | |
| if flag2 == 1: | |
| num_white_pixels = np.sum(image[i + 1] == 255) | |
| if (num_white_pixels > 0.98 * image.shape[1]): | |
| end = i + 1 | |
| break | |
| line = image[int(start - 0.2 * (end - start + 1)): int(end + 1 + 0.2 * (end - start + 1))][:] | |
| if line.shape[0] > 20: | |
| line_rgb = cv.cvtColor(line, cv.COLOR_GRAY2RGB) | |
| lines.append(line_rgb) | |
| pads = 255 * np.ones((20, image.shape[1]), dtype='uint8') | |
| new_image = image[int(end + 2 -(0.2 * (end - start + 1))):][:] | |
| new_image = np.concatenate((pads, new_image)) | |
| image = new_image | |
| return lines | |
| lines = split_image_into_lines(bin_image) | |
| def generate_text(line): | |
| pixel_values = processor(images=line, return_tensors="pt").pixel_values | |
| generated_ids = model.generate(pixel_values, max_new_tokens=50) | |
| generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| return generated_text | |
| with ProcessPoolExecutor() as executor: | |
| results = ' '.join(executor.map(generate_text, lines)) | |
| #improve results with llm | |
| client = OpenAI() | |
| completion = client.chat.completions.create( | |
| model="gpt-4o", | |
| messages=[ | |
| { | |
| "role": "user", | |
| "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}" | |
| } | |
| ] | |
| ) | |
| improved_text = completion.choices[0].message.content | |
| def put_text(text, font, font_scale, color, thickness, max_width, out_image_width, top_margin): | |
| words = text.split(" ") | |
| lines = [] | |
| current_line = "" | |
| for word in words: | |
| if cv.getTextSize(current_line + " " + word, font, font_scale, thickness)[0][0] <= (max_width * out_image_width): | |
| current_line += " " + word | |
| else: | |
| lines.append(current_line) | |
| current_line = word | |
| lines.append(current_line) | |
| 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 | |
| out_image = 255 * (np.ones((out_image_height, out_image_width, 3), dtype=np.uint8)) | |
| top = top_margin | |
| for line in lines: | |
| cv.putText(out_image, line.strip(), (int(((1 - max_width) * out_image_width) / 2), top), font, font_scale, 0, thickness, lineType=cv.LINE_AA) | |
| top += cv.getTextSize(line.strip(), font, font_scale, thickness)[0][1] + 20 | |
| return out_image | |
| font = cv.FONT_HERSHEY_DUPLEX | |
| font_scale = 2 | |
| color = 0 | |
| thickness = 2 | |
| max_width = 0.9 | |
| out_image_width = 1500 | |
| top_margin = 100 | |
| out_image = put_text(improved_text, font, font_scale, color, thickness, max_width, out_image_width, top_margin) |