Update passport_img.py

passport_img.py

@@ -1,6 +1,12 @@
 import cv2
 import numpy as np
 import matplotlib.pyplot as plt
+import ultralytics
+from ultralytics import YOLO
+print(ultralytics.__version__)
+from ultralytics.yolo.utils.ops import scale_image
+model = YOLO('yolov8n-face.pt')
+model_seg = YOLO('yolov8s-seg.pt')
 
 
 
@@ -94,4 +100,154 @@ def merge_all(number,input_path,size):
     # print(merge_img_final.shape)
     # cv2.imwrite(output_path, merge_img_final)
     # print(merge_img_final.shape)
-    return merge_img_final
+    return merge_img_final
+
+def align_crop_image(img):
+#     img = cv2.imread(img)
+    img = cv2.copyMakeBorder(img, 50, 0, 0, 0, cv2.BORDER_CONSTANT, value=(255,255,255))
+    results = model(img)
+    if len(results[0].boxes.data)==1:
+        b_box_up = int(results[0].boxes.data[0][0]),int(results[0].boxes.data[0][1])
+        b_box_down = int(results[0].boxes.data[0][2]),int(results[0].boxes.data[0][3])
+        left_eye=int(results[0].keypoints.data[0][0][0]),int(results[0].keypoints.data[0][0][1])
+        right_eye=int(results[0].keypoints.data[0][1][0]),int(results[0].keypoints.data[0][1][1])
+        nose=int(results[0].keypoints.data[0][2][0]),int(results[0].keypoints.data[0][2][1])
+        left_lip=int(results[0].keypoints.data[0][3][0]),int(results[0].keypoints.data[0][3][1])
+        right_lip=int(results[0].keypoints.data[0][4][0]),int(results[0].keypoints.data[0][4][1])
+        left_eye_center = left_eye
+        left_eye_x = left_eye_center[0] 
+        left_eye_y = left_eye_center[1]
+
+        right_eye_center = right_eye
+        right_eye_x = right_eye_center[0]
+        right_eye_y = right_eye_center[1]
+        
+        if left_eye_y > right_eye_y:
+            A = (right_eye_x, left_eye_y)
+            # Integer -1 indicates that the image will rotate in the clockwise direction
+            direction = -1 
+        else:
+            A = (left_eye_x, right_eye_y)
+            # Integer 1 indicates that image will rotate in the counter clockwise  
+            # direction
+            direction = 1 
+        delta_x = right_eye_x - left_eye_x
+        delta_y = right_eye_y - left_eye_y
+        angle=np.arctan(delta_y/delta_x)
+        angle = (angle * 180) / np.pi
+        h, w = img.shape[:2]
+        # Calculating a center point of the image
+        # Integer division "//"" ensures that we receive whole numbers
+        center = (w // 2, h // 2)
+        # Defining a matrix M and calling
+        # cv2.getRotationMatrix2D method
+        M = cv2.getRotationMatrix2D(center, (angle), 1.0)
+        # Applying the rotation to our image using the
+        # cv2.warpAffine method
+        rotated = cv2.warpAffine(img, M, (w, h))
+        results = model(rotated)
+        b_box_up = int(results[0].boxes.data[0][0]),int(results[0].boxes.data[0][1])
+        b_box_down = int(results[0].boxes.data[0][2]),int(results[0].boxes.data[0][3])
+        box_height = b_box_down[0]-b_box_up[0]
+        box_width = b_box_down[1]-b_box_up[1]
+        left_eye=int(results[0].keypoints.data[0][0][0]),int(results[0].keypoints.data[0][0][1])
+        right_eye=int(results[0].keypoints.data[0][1][0]),int(results[0].keypoints.data[0][1][1])
+        nose=int(results[0].keypoints.data[0][2][0]),int(results[0].keypoints.data[0][2][1])
+        left_lip=int(results[0].keypoints.data[0][3][0]),int(results[0].keypoints.data[0][3][1])
+        right_lip=int(results[0].keypoints.data[0][4][0]),int(results[0].keypoints.data[0][4][1])
+        left_eye_center = left_eye
+        left_eye_x = left_eye_center[0] 
+        left_eye_y = left_eye_center[1]
+
+        right_eye_center = right_eye
+        right_eye_x = right_eye_center[0]
+        right_eye_y = right_eye_center[1]
+
+        up_crop = b_box_up[0]-int(box_height/2),b_box_up[1]-int(box_width/3)
+        down_crop = b_box_down[0]+int(box_height/2),right_lip[1]+box_width
+        final_image = rotated[up_crop[1]:down_crop[1],up_crop[0]:down_crop[0]]
+        resized_image =cv2.resize(final_image, (210, 297))
+        return resized_image
+    else:
+        return None
+    
+    
+def predict_on_image(img, conf):
+    result = model_seg(img, conf=conf)[0]
+
+    # detection
+    # result.boxes.xyxy   # box with xyxy format, (N, 4)
+    cls = result.boxes.cls.cpu().numpy()    # cls, (N, 1)
+    probs = result.boxes.conf.cpu().numpy()  # confidence score, (N, 1)
+    boxes = result.boxes.xyxy.cpu().numpy()   # box with xyxy format, (N, 4)
+
+    # segmentation
+    if result.masks:
+        masks = result.masks.data.cpu().numpy()     # masks, (N, H, W)
+        masks = np.moveaxis(masks, 0, -1) # masks, (H, W, N)
+        print(masks.shape)
+        # rescale masks to original image
+        masks = scale_image(masks, result.masks.orig_shape)
+        masks = np.moveaxis(masks, -1, 0) # masks, (N, H, W)
+
+        return boxes, masks, cls, probs
+    else:
+        return None
+
+
+def overlay(image, mask, color, alpha, resize=None):
+    """Combines image and its segmentation mask into a single image.
+    https://www.kaggle.com/code/purplejester/showing-samples-with-segmentation-mask-overlay
+
+    Params:
+        image: Training image. np.ndarray,
+        mask: Segmentation mask. np.ndarray,
+        color: Color for segmentation mask rendering.  tuple[int, int, int] = (255, 0, 0)
+        alpha: Segmentation mask's transparency. float = 0.5,
+        resize: If provided, both image and its mask are resized before blending them together.
+        tuple[int, int] = (1024, 1024))
+
+    Returns:
+        image_combined: The combined image. np.ndarray
+
+    """
+    for i in range(len(mask)):
+        for j in range(len(mask[i])):
+    #         print(masks[0][i][j])
+            mask[i][j] = abs(mask[i][j] - 1)
+    color = color[::-1]
+    colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
+    colored_mask = np.moveaxis(colored_mask, 0, -1)
+    masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
+    image_overlay = masked.filled()
+    if resize is not None:
+        image = cv2.resize(image.transpose(1, 2, 0), resize)
+        image_overlay = cv2.resize(image_overlay.transpose(1, 2, 0), resize)
+
+    image_combined = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)
+
+    return image_combined
+
+def add_name_dob(img,name,dob):
+    img = cv2.resize(img,(420,594))
+    img = cv2.rectangle(img,(0,520),(420,594),(255,255,255),-1)
+    dob = 'D.O.B: '+dob
+    font = 1#cv2.FONT_HERSHEY_SIMPLEX
+    font_scale = 2
+    font_thickness = 2
+
+    # Get the size of the text
+    text_size = cv2.getTextSize(name, font, font_scale, font_thickness)[0]
+    dob_size = cv2.getTextSize(dob, font, font_scale, font_thickness)[0]
+
+    # Calculate the position to center the text
+    text_x = (img.shape[1] - text_size[0]) // 2
+    text_y = (img.shape[0] + text_size[1]) // 2
+    dob_x = (img.shape[1] - dob_size[0]) // 2
+    dob_y = (img.shape[0] + dob_size[1]) // 2
+
+    # Put the text on the image
+    cv2.putText(img, name, (text_x, 547), font, font_scale, (0, 0, 0), font_thickness)
+    cv2.putText(img, dob, (dob_x, 583), font, font_scale, (0, 0, 0), font_thickness)
+
+    return img