CadExtractor / src /utils /imageFormater.py
Martin Krockert
Demo with tesseract / paddle and finetuned yolo 12
fa54254
raw
history blame
1.77 kB
import PIL.Image
import numpy as nd
import cv2
import io
import base64
from io import BytesIO
def toPIL(array: nd.array) -> PIL.Image:
return PIL.Image.fromarray(array).convert('RGB')
def toBytes(pilImage : PIL.Image) -> bytes:
img_byte_arr = io.BytesIO()
pilImage.save(img_byte_arr, format='JPEG')
return img_byte_arr.getvalue()
def toBASE64(pilImage : PIL.Image) -> str:
buffered = BytesIO()
pilImage.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue())
def redLines(bgr) -> nd.array:
# Convert the image to grayscale#
if len(bgr) == 0:
return bgr
cont = bgr.copy()
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
# Threshold the image to create a binary mask
_, thresh = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)
# Set the line size (adjust as needed)
linesize = 25
# Remove horizontal lines
horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (linesize, 1))
remove_horizontal = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
cnts = cv2.findContours(remove_horizontal, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
for c in cnts:
cv2.drawContours(cont, [c], -1, (255, 0, 0), 2) # Draw in red
# Remove vertical lines
vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, linesize))
remove_vertical = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
cnts = cv2.findContours(remove_vertical, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
for c in cnts:
cv2.drawContours(cont, [c], -1, (255, 0, 0), 2) # Draw in red
return cont