import re import streamlit as st from PIL import Image from PIL import Image import os import argparse import torch import torch.nn as nn import torch.backends.cudnn as cudnn from torch.autograd import Variable from test import copyStateDict from PIL import Image import cv2 from skimage import io import numpy as np import test import file_utils import pandas as pd from craft import CRAFT from collections import OrderedDict from PIL import Image from vietocr.tool.predictor import Predictor from vietocr.tool.config import Cfg import os import matplotlib.pyplot as plt import numpy as np from pathlib import Path import cv2 import glob import matplotlib.pyplot as plt from PIL import Image, ImageTk import file_utils import os import imgproc # Tạo yêu cầu đến mô hình def str2bool(v): return v.lower() in ("yes", "y", "true", "t", "1") ''' ''' # CRAFT parser = argparse.ArgumentParser(description='CRAFT Text Detection') parser.add_argument('--trained_model', default='weights/craft_mlt_25k.pth', type=str, help='pretrained model') parser.add_argument('--text_threshold', default=0.7, type=float, help='text confidence threshold') parser.add_argument('--low_text', default=0.4, type=float, help='text low-bound score') parser.add_argument('--link_threshold', default=0.4, type=float, help='link confidence threshold') parser.add_argument('--cpu', default=True, type=str2bool, help='Use cpu for inference') parser.add_argument('--canvas_size', default=1280, type=int, help='image size for inference') parser.add_argument('--mag_ratio', default=1.5, type=float, help='image magnification ratio') parser.add_argument('--poly', default=False, action='store_true', help='enable polygon type') parser.add_argument('--show_time', default=False, action='store_true', help='show processing time') parser.add_argument('--test_folder', default='data_image', type=str, help='đường dẫn tới ảnh đầu vào') parser.add_argument('--refine', default=True, action='store_true', help='enable link refiner') parser.add_argument('--refiner_model', default='weights/craft_refiner_CTW1500.pth', type=str, help='pretrained refiner model') args = parser.parse_args() ######################################################################################### csv_columns = ['x_top_left', 'y_top_left', 'x_top_right', 'y_top_right', 'x_bot_right', 'y_bot_right', 'x_bot_left', 'y_bot_left'] # load net net = CRAFT() # initialize print('Đang thực hiện load weight (' + args.trained_model + ')') if args.cpu: net.load_state_dict(copyStateDict(torch.load(args.trained_model, map_location='cpu'))) else: net.load_state_dict(copyStateDict(torch.load(args.trained_model, map_location='cpu'))) if args.cpu: net = net.cpu() net = torch.nn.DataParallel(net) cudnn.benchmark = False net.eval() # LinkRefiner Đoạn này code không chạy qua nên không cần đọc vì weight đã load ở cái bên trên # còn refine để mặc định bên trên là False nên sẽ bị bỏ qua # ------------------------------------------------------------------------------------------ # ------------------------------------------------------------------------------------------ refine_net = None if args.refine: from refinenet import RefineNet refine_net = RefineNet() print('Đang thực hiện load weight (' + args.refiner_model + ')') if args.cpu: refine_net.load_state_dict(copyStateDict(torch.load(args.refiner_model, map_location='cpu'))) refine_net = refine_net.cpu() refine_net = torch.nn.DataParallel(refine_net) else: refine_net.load_state_dict(copyStateDict(torch.load(args.refiner_model, map_location='cpu'))) refine_net.eval() args.poly = True config = Cfg.load_config_from_name('vgg_transformer') config['export'] = 'transformerocr_checkpoint.pth' config['device'] = 'cpu' config['predictor']['beamsearch'] = False detector = Predictor(config) # ------------------------------------------------------------------------------------------ # ------------------------------------------------------------------------------------------ # Tạo tiêu đề và phần tải lên hình ảnh st.title("Trích xuất thông tin từ căn cước công dân") uploaded_file = st.file_uploader("Tải lên ảnh căn cước công dân", type=["jpg", "jpeg", "png"]) if uploaded_file is not None: # Hiển thị hình ảnh tải lên image = Image.open(uploaded_file) image.save("uploaded_image.jpg") st.image(image, caption='Hình ảnh căn cước', use_column_width=True) import tempfile # Lưu trữ tạm thời và lấy đường dẫn with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: # Save the image to the temp file temp_file.write(uploaded_file.read()) temp_file_path = temp_file.name # Hiển thị đường dẫn tạm thời print(f"Đường dẫn tạm thời của file: {temp_file_path}") # Nút chạy if st.button("Run"): print("ok") image_path = "uploaded_image.jpg" k = 1 crop_folder = "crop_Word" result_folder = "Results" image = imgproc.loadImage(image_path) bboxes, polys, score_text, det_scores = test.test_net(net, image, args.text_threshold, args.link_threshold, args.low_text, args.cpu, args.poly, args, refine_net) bbox_score = {} def crop_polygon(image, vertices, box_num1): # Tạo mặt nạ mask = np.zeros(image.shape[:2], dtype=np.uint8) cv2.fillPoly(mask, [np.int32(vertices)], 255) # Tìm bounding rect để crop vùng chứa đa giác rect = cv2.boundingRect(np.int32(vertices)) # Crop và lấy hình ảnh con theo bounding rect cropped = image[rect[1]:rect[1]+rect[3], rect[0]:rect[0]+rect[2]] # Tạo mặt nạ cho vùng đã crop cropped_mask = mask[rect[1]:rect[1]+rect[3], rect[0]:rect[0]+rect[2]] # Lọc vùng bằng mặt nạ result = cv2.bitwise_and(cropped, cropped, mask=cropped_mask) crop_path = os.path.join(crop_folder, f"crop_{box_num1 + 1}.jpg") cv2.imwrite(crop_path, result) return result if len(bboxes) == 0: with open(f"data_text//text_{k}.txt", "w", encoding="utf-8") as f: f.write(" ") else: for box_num in range(len(bboxes)): item = bboxes[box_num] data = np.array([[int(item[0][0]), int(item[0][1]), int(item[1][0]), int(item[1][1]), int(item[2][0]), int(item[2][1]), int(item[3][0]), int(item[3][1])]]) csvdata = pd.DataFrame(data, columns=csv_columns) csvdata.to_csv(f'data{k}.csv', index=False, mode='a', header=False) # save score text filename, file_ext = os.path.splitext(os.path.basename(image_path)) mask_file = result_folder + "/res_" + filename + '_mask.jpg' # tạo đường dẫn file bản đồ nhiệt cv2.imwrite(mask_file, score_text) # in ra bản đồ nhiệt # file_utils.saveResult(image_path, image[:, :, ::-1], polys, dirname=result_folder) cropped_images = [] for i, box in enumerate(bboxes): cropped = crop_polygon(image, box, i) cropped_images.append(cropped) print(f"Đã cắt {len(cropped_images)} vùng bounding box.") path = glob.glob("crop_Word/*.jpg") cv_img = [str(detector.predict(Image.open(f'crop_Word/crop_' + str(i + 1) + '.jpg'))) for i in range(len(bboxes))] print(cv_img) from google.generativeai.types import HarmCategory, HarmBlockThreshold import google.generativeai as genai genai.configure(api_key="AIzaSyAH4ayK6nL71wxPtuYOCe32OdZVZAANWic") # Khởi tạo mô hình model = genai.GenerativeModel(model_name='gemini-1.5-flash') # Thiết lập safe_setting cho các loại harm có sẵn và hợp lệ safety_settings = [ {"category": HarmCategory.HARM_CATEGORY_HATE_SPEECH, "threshold": HarmBlockThreshold.BLOCK_NONE}, {"category": HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, "threshold": HarmBlockThreshold.BLOCK_NONE}, {"category": HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, "threshold": HarmBlockThreshold.BLOCK_NONE}, {"category": HarmCategory.HARM_CATEGORY_HARASSMENT, "threshold": HarmBlockThreshold.BLOCK_NONE} ] print("ok") response = model.generate_content( [f"tìm tên, ngày sinh, nơi cư trú, số căn cước, hạn sử dụng trong mảng sau {cv_img}, chỉ trả về tên, ngày sinh, nơi cư trú, số căn cước,hạn sử dụng mà model tìm thấy được, không trả lời thêm gì, ví dụ 'NGUYỄN THANH SANG, 18/05/1981, 223/11 Kv Bỉnh- Dương, Long Hòa, Bình Thủy, Cần Thơ, 092081007131, 18/05/2041 '"], safety_settings=safety_settings ) # print(f"tìm tên, ngày sinh, nơi cư trú, số căn cước, hạn sử dụng trong mảng sau {cv_img}, chỉ trả về tên, ngày sinh, nơi cư trú, số căn cước,hạn sử dụng mà model tìm thấy được, không trả lời thêm gì, ví dụ 'NGUYỄN THANH SANG, 18/05/1981, 223/11 Kv Bỉnh- Dương, Long Hòa, Bình Thủy, Cần Thơ, 092081007131, 18/05/2041 '") print(response.text) # for box in bboxes: # cv2.polylines(image, [np.int32(box)], isClosed=True, color=(0, 255, 0), thickness=1) # plt.figure(figsize=(20, 20)) # plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) # plt.title('Detected Text Bounding Boxes') # plt.show() print(f"đã load xong ảnh {k + 1}") info_string = response.text.strip() # Loại bỏ khoảng trắng thừa print(type(info_string)) print(info_string) # In chuỗi để kiểm tra pattern = r"^(.+?),\s*(\d{2}/\d{2}/\d{4}),\s*(.+?),\s*(.+?),\s*(.+?),\s*(\d{12}),\s*(\d{2}/\d{2}/\d{4})$" match = re.match(pattern, info_string) if match: ho_ten = match.group(1) # ĐỖ VĂN MẠNH ngay_sinh = match.group(2) # 18/10/1991 noi_cu_tru = match.group(3) + ", " + match.group(4) + ", " + match.group(5) # Tân Dân, Thành phố chí Linh, Hải Dương so_can_cuoc = match.group(6) # 030091002288 han_su_dung = match.group(7) # 18/10/2031 # Hiển thị các thông tin đã lấy được print("Họ và tên:", ho_ten) print("Ngày sinh:", ngay_sinh) print("Nơi cư trú:", noi_cu_tru) print("Số căn cước:", so_can_cuoc) print("Hạn sử dụng:", han_su_dung) else: print("Không tìm thấy thông tin phù hợp trong chuỗi.") # Hiển thị kết quả st.subheader("Kết quả trích xuất:") st.text_area("ALL TEXT", cv_img) st.text_area("Thông tin căn cước", response.text) st.text_area("Tên", ho_ten if match else "") st.text_area("Ngày sinh", ngay_sinh if match else "") st.text_area("Nơi cư trú:", noi_cu_tru if match else "") st.text_area("Số căn cước", so_can_cuoc if match else "") st.text_area("Hạn sử dụng", han_su_dung if match else "")