Spaces:
Sleeping
Sleeping
| import random | |
| from PIL import ImageFont, ImageDraw, Image | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import src.solve as solve | |
| from typing import Tuple | |
| import pytesseract | |
| import re | |
| import shutil | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| import random | |
| # Setup class names | |
| with open("class_names.txt", "r") as f: # reading them in from class_names.txt | |
| class_names = [names.strip() for names in f.readlines()] | |
| model1 = tf.keras.models.load_model('model/model30.h5') | |
| model2 = tf.keras.models.load_model('model/model15.h5') | |
| model3 = tf.keras.models.load_model('model/model2.h5') | |
| palabras_1 = [] | |
| # Borrar el directorio de imagenes | |
| folder = 'output' | |
| for filename in os.listdir(folder): | |
| file_path = os.path.join(folder, filename) | |
| try: | |
| if os.path.isfile(file_path) or os.path.islink(file_path): | |
| os.unlink(file_path) | |
| elif os.path.isdir(file_path): | |
| shutil.rmtree(file_path) | |
| except Exception as e: | |
| print('Failed to delete %s. Reason: %s' % (file_path, e)) | |
| folder = 'wordsPuzzle' | |
| for filename in os.listdir(folder): | |
| file_path = os.path.join(folder, filename) | |
| try: | |
| if os.path.isfile(file_path) or os.path.islink(file_path): | |
| os.unlink(file_path) | |
| elif os.path.isdir(file_path): | |
| shutil.rmtree(file_path) | |
| except Exception as e: | |
| print('Failed to delete %s. Reason: %s' % (file_path, e)) | |
| def get_words(img): | |
| #print(type(img)) | |
| # str to filepath | |
| img = Image.open(img) | |
| # Display image | |
| # img.show() | |
| text = pytesseract.image_to_string(img, lang="spa+eng", config="--psm 11") | |
| text = text.upper() | |
| text = re.split('\W+', text) | |
| text.pop() | |
| #palabras_1 = text | |
| #print(palabras_1) | |
| # array to string text | |
| text = ' '.join(text) | |
| # add comma to text | |
| text = text.replace(" ", ",") | |
| return text | |
| def getmat(listaCuadrados, filas, columnas): | |
| matrix = [[0 for i in range(columnas)] for j in range(filas)] | |
| matrixT = [[0 for i in range(columnas)] for j in range(filas)] | |
| listaCuadrados.sort(key=lambda y: y["altura"]) | |
| e = 0 | |
| for i in range(filas): | |
| lista2 = listaCuadrados[e:e+columnas] | |
| e = e+columnas | |
| lista2.sort(key=lambda y: y["anchura"]) | |
| j = 0 | |
| for lista in lista2: | |
| matrix[i][j] = lista["letra"] | |
| matrixT[i][j] = lista | |
| j = j + 1 | |
| return matrix, matrixT | |
| # Obtener las filas y las columnas de la sopa | |
| def get_colums_and_rows(listaCuadrados): | |
| columnas = 0 | |
| filas = 1 | |
| alturaAnt = listaCuadrados[0]["altura"] | |
| for lista in listaCuadrados: | |
| if (lista["altura"] > alturaAnt + 6): | |
| filas = filas + 1 | |
| alturaAnt = lista["altura"] | |
| if (filas == 1): | |
| columnas = columnas + 1 | |
| return filas, columnas | |
| def read_board(img, words): | |
| #(type(img)) | |
| # str to filepath | |
| img = Image.open(img) | |
| # Display image | |
| img.show() | |
| # Print words | |
| #print("Palabras a buscar: ", palabras_1) | |
| def solve_puzzle(img, words): | |
| # Borrar el directorio de imagenes | |
| folder = 'output' | |
| for filename in os.listdir(folder): | |
| file_path = os.path.join(folder, filename) | |
| try: | |
| if os.path.isfile(file_path) or os.path.islink(file_path): | |
| os.unlink(file_path) | |
| elif os.path.isdir(file_path): | |
| shutil.rmtree(file_path) | |
| except Exception as e: | |
| print('Failed to delete %s. Reason: %s' % (file_path, e)) | |
| folder = 'wordsPuzzle' | |
| for filename in os.listdir(folder): | |
| file_path = os.path.join(folder, filename) | |
| try: | |
| if os.path.isfile(file_path) or os.path.islink(file_path): | |
| os.unlink(file_path) | |
| elif os.path.isdir(file_path): | |
| shutil.rmtree(file_path) | |
| except Exception as e: | |
| print('Failed to delete %s. Reason: %s' % (file_path, e)) | |
| # print(type(img)) | |
| # str to filepath | |
| #print(type(words)) | |
| #print(words) | |
| # img = Image.open(img) | |
| # Pil to opencv compatible | |
| pil_image = Image.open(img).convert('RGB') | |
| open_cv_image = np.array(pil_image) | |
| # Convert RGB to BGR | |
| open_cv_image = open_cv_image[:, :, ::-1].copy() | |
| # Display image | |
| # Print words | |
| img = open_cv_image | |
| # string to array | |
| words = words.split(",") | |
| # remove last , | |
| #print(words) | |
| imgc = img.copy() | |
| imgsol = img.copy() | |
| imgc = cv2.cvtColor(imgc, cv2.COLOR_RGB2GRAY) | |
| imgc = np.invert(imgc) | |
| gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) | |
| blur = cv2.GaussianBlur(gray, (5, 5), 0) | |
| # save the blurred image | |
| #cv2.imwrite("output/blur.png", blur) | |
| # display blurred image | |
| threshten = cv2.threshold( | |
| blur, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] | |
| thresh = cv2.adaptiveThreshold(threshten, 255, 1, 1, 11, 2) | |
| contours, hierarchy = cv2.findContours( | |
| thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| # Draw contours and save the image | |
| characters = np.array([ | |
| 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', | |
| 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', | |
| 'U', 'V', 'W', 'X', 'Y', 'Z', 'Á', 'É', 'Í', 'Ñ', 'Ó', 'Ú', '?']) | |
| offset = 7 | |
| minh = 7 | |
| minw = 2 | |
| i = 0 | |
| listaCuadrados = list() | |
| fnt = ImageFont.truetype("fonts/Roboto-Black.ttf", 20) | |
| for cnt in reversed(contours): | |
| contCuadrados = {} | |
| if cv2.contourArea(cnt) > 18 and cv2.contourArea(cnt) < 1000: | |
| [x, y, w, h] = cv2.boundingRect(cnt) | |
| if h > minh and w > minw: | |
| i = i + 1 | |
| height, width = imgc.shape | |
| y0 = 0 | |
| y1 = height | |
| x0 = 0 | |
| x1 = width | |
| if height > y+h+offset: | |
| y1 = y+h+offset | |
| if width > x+w+offset: | |
| x1 = x+w+offset | |
| if y-offset > 0: | |
| y0 = y-offset | |
| if x-offset > 0: | |
| x0 = x-offset | |
| cv2.rectangle(img, (x0, y0), (x1, y1), (0, 255, 0), 2) | |
| img2 = imgc[y0:y1, x0:x1] | |
| img2 = cv2.resize(img2, (28, 28)) | |
| img_array = img2.reshape(1, 28, 28, 1) | |
| prediction1 = np.argmax(model1.predict(img_array)) | |
| prediction2 = np.argmax(model2.predict(img_array)) | |
| prediction3 = np.argmax(model3.predict(img_array)) | |
| pred = 0 | |
| if prediction1 == prediction2 and prediction2 == prediction3: | |
| pred = prediction1 | |
| elif prediction1 == prediction2: | |
| pred = prediction1 | |
| elif prediction2 == prediction3: | |
| pred = prediction2 | |
| elif prediction1 == prediction3: | |
| pred = prediction3 | |
| else: | |
| pred = 32 | |
| #print(characters[pred]) | |
| contCuadrados["anchura"] = x | |
| contCuadrados["altura"] = y | |
| contCuadrados["centrox"] = (x + x + w)/2 | |
| contCuadrados["centroy"] = (y + y + h)/2 | |
| contCuadrados["letra"] = characters[pred] | |
| listaCuadrados.append(contCuadrados) | |
| img_pil = Image.fromarray(img) | |
| draw = ImageDraw.Draw(img_pil) | |
| draw.text(((x0+x1)/2, y0-10), | |
| characters[pred], font=fnt, fill=(255, 0, 0, 0)) | |
| img = np.array(img_pil) | |
| #cv2.imwrite("output/Tablero_Labels.png", img) | |
| filas, columnas = get_colums_and_rows(listaCuadrados) | |
| # print listaCuadrados | |
| # print(listaCuadrados) | |
| matrix, matrixT = getmat(listaCuadrados, filas, columnas) | |
| palabrasxy = [] | |
| # print() | |
| # print("Palabras a buscar:") | |
| # for i in words: | |
| # print(i) | |
| # print() | |
| # for i in range(filas): | |
| # for j in range(columnas): | |
| # print(matrix[i][j], end = " ") | |
| # print() | |
| # print() | |
| # print() | |
| image_new = imgsol.copy() | |
| overlay = imgsol.copy() | |
| index = 0 | |
| index2 = 0 | |
| for i in words: | |
| #(i) | |
| xy_positionsvec, find = solve.find_word(matrix, i) | |
| if find: | |
| palabrasxy.append(xy_positionsvec) | |
| xy_positionsvec = palabrasxy[index] | |
| # print(len(xy_positionsvec)) | |
| xy = xy_positionsvec[0] | |
| xy2 = xy_positionsvec[len(xy_positionsvec)-1] | |
| # print(xy["x"], " ",xy["y"]) | |
| # print(xy2["x"], " ",xy2["y"]) | |
| coordreal = matrixT[xy["x"]][xy["y"]] | |
| coordreal2 = matrixT[xy2["x"]][xy2["y"]] | |
| centrox = coordreal["centrox"] | |
| centroy = coordreal["centroy"] | |
| centrox2 = coordreal2["centrox"] | |
| centroy2 = coordreal2["centroy"] | |
| centrox = int(centrox) | |
| centroy = int(centroy) | |
| centrox2 = int(centrox2) | |
| centroy2 = int(centroy2) | |
| color = (random.randint(0, 255), random.randint( | |
| 0, 255), random.randint(0, 255)) | |
| cv2.line(overlay, (centrox, centroy), (centrox2, centroy2), color, | |
| thickness=int(abs(coordreal["altura"] - coordreal["centroy"])*2)) | |
| overlay2 = imgsol.copy() | |
| cv2.line(overlay2, (centrox, centroy), (centrox2, centroy2), color, | |
| thickness=int(abs(coordreal["altura"] - coordreal["centroy"])*2)) | |
| image_word = cv2.addWeighted(overlay2, 0.4, image_new, 1 - 0.4, 0) | |
| cv2.imwrite("wordsPuzzle/" + words[index2] + ".jpg", image_word) | |
| # append the image into a numpy array | |
| #print(words[index2]) | |
| index += 1 | |
| index2 += 1 | |
| alpha = 0.4 # Transparency factor | |
| image_new = cv2.addWeighted(overlay, alpha, image_new, 1 - alpha, 0) | |
| cv2.imwrite("output/Tablero_solucion.png", image_new) | |
| # return the images in wordsPuzzle folder as numpy arrays | |
| return image_new |