code

Segmentation.py

@@ -0,0 +1,58 @@
+from ultralytics import YOLO
+import cv2
+import os
+import shutil
+import numpy as np
+def black_mask(results):
+    # Đường dẫn đến thư mục "runs"
+    runs_folder = "mask"
+# Kiểm tra xem thư mục "mask" đã tồn tại chưa
+    if not os.path.exists(runs_folder):
+        # Nếu chưa tồn tại, tạo thư mục "mask"
+        os.makedirs(runs_folder)
+    for mask_index in range(len(results[0].masks)):
+        # Lấy mặt nạ
+            mask_raw = results[0].masks[mask_index].cpu().data.numpy().transpose(1, 2, 0)
+            # Chuyển đổi mặt nạ thành 3 kênh màu
+            mask_3channel = cv2.merge((mask_raw, mask_raw, mask_raw))
+            # Lấy kích thước của ảnh gốc (chiều cao, chiều rộng, số kênh)
+            h2, w2, c2 = results[0].orig_img.shape
+            # Thay đổi kích thước mặt nạ thành cùng kích thước với ảnh gốc
+            mask = cv2.resize(mask_3channel, (w2, h2))
+            # Chuyển đổi ảnh mặt nạ thành không gian màu HSV
+            hsv = cv2.cvtColor(mask, cv2.COLOR_BGR2HSV)
+            # Xác định phạm vi độ sáng trong không gian màu HSV
+            lower_black = np.array([0, 0, 0])
+            upper_black = np.array([0, 0, 1])
+            # Tạo mặt nạ dựa trên phạm vi độ sáng
+            mask = cv2.inRange(mask, lower_black, upper_black)
+            # Đảo ngược mặt nạ để lấy mọi thứ trừ màu đen
+            mask = cv2.bitwise_not(mask)
+            # Áp dụng mặt nạ vào ảnh gốc
+            masked = cv2.bitwise_and(results[0].orig_img, results[0].orig_img, mask=mask)
+            # Chuyển mọi pixel màu đen thành màu trắng
+            black_color = [0, 0, 0]
+            white_color = [255, 255, 255]
+            result_image = np.where(np.all(masked == black_color, axis=-1, keepdims=True), masked, white_color)
+            name = "mask/"+str(mask_index) + '.png'
+            # Show the masked part of the image
+            cv2.imwrite(name,result_image)
+def segmentation_doc(image):
+    model = YOLO('Model/Seg65ep.pt')
+    runs_folder = "runs"
+    folder="mask"
+    # Xóa thư mục runs
+    if os.path.exists(runs_folder) and os.path.isdir(runs_folder):
+        # Xóa thư mục runs
+        shutil.rmtree(runs_folder)
+    if os.path.exists(folder) and os.path.isdir(folder):
+        # Xóa thư mục runs
+        shutil.rmtree(folder)
+    color = [255, 255, 255] # Màu trắng
+    padding_size = 50
+    image = cv2.copyMakeBorder(image, padding_size, padding_size, padding_size, padding_size, cv2.BORDER_CONSTANT, value=color)
+    results = model(image,save=True,retina_masks = True, conf=0.7)
+    black_mask(results)
+    return image
+# image=cv2.imread("test1.jpg")
+# segmentation_doc(image)

perpestive_transform.py

@@ -0,0 +1,188 @@
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+from deskew import determine_skew
+import math
+from typing import Tuple, Union
+def rotate(image: np.ndarray, angle: float, background: Union[int, Tuple[int, int, int]]) -> np.ndarray:
+    old_width, old_height = image.shape[:2]
+    angle_radian = math.radians(angle)
+    width = abs(np.sin(angle_radian) * old_height) + abs(np.cos(angle_radian) * old_width)
+    height = abs(np.sin(angle_radian) * old_width) + abs(np.cos(angle_radian) * old_height)
+
+    image_center = tuple(np.array(image.shape[1::-1]) / 2)
+    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
+    rot_mat[1, 2] += (width - old_width) / 2
+    rot_mat[0, 2] += (height - old_height) / 2
+
+    return cv2.warpAffine(image, rot_mat, (int(round(height)), int(round(width))), borderValue=background)
+
+
+def check(my_list):
+    unique_elements = []
+
+    # Sử dụng vòng lặp để kiểm tra từng phần tử trong danh sách
+    for item in my_list:
+        # Nếu phần tử không xuất hiện trong danh sách các phần tử duy nhất, thêm nó vào danh sách đó
+        if item not in unique_elements:
+            unique_elements.append(item)
+    return len(unique_elements)
+
+
+def order_points(pts):
+    rect = np.zeros((4, 2), dtype="float32")
+    pts = np.array(pts)
+    s = pts.sum(axis=1)
+    rect[0] = pts[np.argmin(s)]
+    rect[2] = pts[np.argmax(s)]
+
+    diff = np.diff(pts, axis=1)
+    rect[1] = pts[np.argmin(diff)]
+    rect[3] = pts[np.argmax(diff)]
+    return rect.astype("int").tolist()
+
+
+def find_dest(pts):
+    (tl, tr, br, bl) = pts
+    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
+    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
+    maxWidth = max(int(widthA), int(widthB))
+    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
+    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
+    maxHeight = max(int(heightA), int(heightB))
+    destination_corners = [[0, 0], [maxWidth, 0], [maxWidth, maxHeight], [0, maxHeight]]
+    return order_points(destination_corners)
+def extract(ori_img, img, image_size=384, BUFFER=100):
+    img=img.astype(np.uint8)
+    gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    angle = determine_skew(gray_image)
+    img = rotate(img, angle, (0, 0, 0))
+    ori_img = rotate(ori_img, angle, (0, 0, 0))
+    #get size of image
+    size = img.shape
+    top_pad = 10     # Số pixel padding ở phía trên
+    bottom_pad = 10   # Số pixel padding ở phía dưới
+    left_pad = 10   # Số pixel padding ở phía trái
+    right_pad = 10   # Số pixel padding ở phía phải
+
+    # Tạo hình ảnh mới với kích thước lớn hơn, bằng cách thêm pixel màu đen (0) vào xung quanh
+    height, width, channels = img.shape
+    new_height = height + top_pad + bottom_pad
+    new_width = width + left_pad + right_pad
+
+    # Tạo một hình ảnh mới với màu đen (0) là màu nền
+    padded_img = np.zeros((new_height, new_width, channels), dtype=np.uint8)
+
+    # Copy nội dung của hình ảnh gốc vào vị trí tương ứng trong hình ảnh mới
+    padded_img[top_pad:top_pad + height, left_pad:left_pad + width] = img
+    img = padded_img
+
+    height, width, channels = ori_img.shape
+    new_height = height + top_pad + bottom_pad
+    new_width = width + left_pad + right_pad
+
+    # Tạo một hình ảnh mới với màu đen (0) là màu nền
+    padded_ori_img = np.full((new_height, new_width, channels), 255, dtype=np.uint8)
+
+    # Copy nội dung của hình ảnh gốc vào vị trí tương ứng trong hình ảnh mới
+    padded_ori_img[top_pad:top_pad + height, left_pad:left_pad + width] = ori_img
+    ori_img = padded_ori_img
+
+    imH, imW, C = img.shape
+    IMAGE_SIZE = image_size
+    img_rs = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_NEAREST)
+
+    # imH, imW, C = img.shape
+    # IMAGE_SIZE=image_size
+    scale_x = imW / IMAGE_SIZE
+    scale_y = imH / IMAGE_SIZE
+    # img=cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_NEAREST)
+
+    canny = cv2.Canny(img_rs.astype(np.uint8), 225, 255)
+    canny = cv2.dilate(canny, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)))
+    contours, _ = cv2.findContours(canny, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
+    page = sorted(contours, key=cv2.contourArea, reverse=True)[0]
+    epsilon = (0.02* cv2.arcLength(page, True))
+    corners = cv2.approxPolyDP(page, epsilon, True)
+    corners = np.concatenate(corners).astype(np.float32)
+    corners[:, 0] *= scale_x
+    corners[:, 1] *= scale_y
+
+    # corners[:, 0] -= half
+    # corners[:, 1] -= half
+    for corner in corners:
+        x, y = corner.astype(int)
+        cv2.circle(img, (int(x), int(y)), 20, (0, 255, 0), -1)  # Vẽ một hình tròn màu xanh lên ảnh
+
+    if len(corners) > 4:
+        left_pad, top_pad, right_pad, bottom_pad = 0, 0, 0, 0
+        rect = cv2.minAreaRect(corners.reshape((-1, 1, 2)))
+        box = cv2.boxPoints(rect)
+        box_corners = np.int32(box)
+        #     box_corners = minimum_bounding_rectangle(corners)
+
+        box_x_min = np.min(box_corners[:, 0])
+        box_x_max = np.max(box_corners[:, 0])
+        box_y_min = np.min(box_corners[:, 1])
+        box_y_max = np.max(box_corners[:, 1])
+
+        # Find corner point which doesn't satify the image constraint
+        # and record the amount of shift required to make the box
+        # corner satisfy the constraint
+        if box_x_min <= 0:
+            left_pad = abs(box_x_min) + BUFFER
+
+        if box_x_max >= imW:
+            right_pad = (box_x_max - imW) + BUFFER
+
+        if box_y_min <= 0:
+            top_pad = abs(box_y_min) + BUFFER
+
+        if box_y_max >= imH:
+            bottom_pad = (box_y_max - imH) + BUFFER
+        box_corners[:, 0] += left_pad
+        box_corners[:, 1] += top_pad
+        corners = box_corners
+    
+    if check(order_points(corners)) >= 4:
+        corners = order_points(corners)
+    else:
+        pass
+
+# Define the amount to increase the rectangle size
+    # (tl, tr, br, bl) = corners
+    # widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
+    # widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
+    # maxWidth = max(int(widthA), int(widthB))
+    # heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
+    # heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
+    # maxHeight = max(int(heightA), int(heightB))
+
+    # # Increase the x-coordinate of the top-right and bottom-right points
+    # corners[1][0] += maxWidth/30
+    # corners[2][0] += maxWidth/30
+
+    # # Decrease the x-coordinate of the top-left and bottom-left points
+    # corners[0][0] -= maxWidth/30
+    # corners[3][0] -= maxWidth/30
+    # # Increase the y-coordinate of the bottom-right and bottom-left points
+    # corners[2][1] += maxHeight/30
+    # corners[3][1] += maxHeight/30
+
+    # # Decrease the y-coordinate of the top-left and top-right points
+    # corners[0][1] -= maxHeight/30
+    # corners[1][1] -= maxHeight/30
+
+    # print(corners)
+
+    destination_corners = find_dest(corners)
+    M = cv2.getPerspectiveTransform(np.float32(corners), np.float32(destination_corners))
+    final = cv2.warpPerspective(ori_img, M, (destination_corners[2][0], destination_corners[2][1]),
+                                flags=cv2.INTER_LANCZOS4)
+    return final
+# ori=cv2.imread("runs\segment\predict\image0.jpg")
+# img=cv2.imread("mask/2.png")
+# final=extract(ori,img)
+# plt.imshow(final)
+# plt.show()
+# # print(img.shape)