Spaces:
Running
Running
| 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') | |
| def crop_passport_size(img_arr,size): | |
| if size==3545: | |
| passport_size = (310, 410) | |
| elif size==34: | |
| passport_size = (260, 346) | |
| else: | |
| passport_size = (170, 210) | |
| # Load the input image | |
| # original_image = cv2.imread(input_image_path) | |
| original_image = img_arr | |
| # Define the dimensions for the passport-size photo (35mm x 45mm) | |
| # passport_size = (310, 410) | |
| # Resize the image to the passport-size dimensions | |
| resized_image = cv2.resize(original_image, (passport_size[0], passport_size[1])) | |
| # Save the cropped passport-size image | |
| # w_output = add_border(output, 4, 4, 4, 4, (0,0,0)) | |
| # w_output = add_border(w_output, 15, 15, 15, 15, (255,255,255)) | |
| # cv2.imwrite(output_image_path, resized_image) | |
| return resized_image | |
| def add_border(image, top, bottom, left, right, color): | |
| return cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) | |
| def horizontal_merge(number,crop_img,size): | |
| if size==3545: | |
| white_image = np.ones((450, 350, 3), dtype=np.uint8) * 255 | |
| img_per_row = 6 | |
| elif size==34: | |
| white_image = np.ones((386, 300, 3), dtype=np.uint8) * 255 | |
| img_per_row = 7 | |
| else: | |
| white_image = np.ones((250, 210, 3), dtype=np.uint8) * 255 | |
| img_per_row = 10 | |
| # white_image = np.ones((450, 350, 3), dtype=np.uint8) * 255 | |
| list_img = [] | |
| # for i in range(5): | |
| for i in range(number): | |
| list_img.append(crop_img) | |
| for i in range(img_per_row-number): | |
| list_img.append(white_image) | |
| merge_img = cv2.hconcat(list_img) | |
| return merge_img | |
| def merge_all(number,input_path,size): | |
| input_image_path = input_path # Replace with the path to your input image | |
| # output_image_path = output_path # Replace with the desired output path | |
| output=crop_passport_size(input_image_path,size) | |
| w_output = add_border(output, 4, 4, 4, 4, (0,0,0)) | |
| w_output = add_border(w_output, 15, 15, 15, 15, (255,255,255)) | |
| w_output = add_border(w_output, 1, 1, 1, 1, (0,0,0)) | |
| if size==3545: | |
| img_per_row = 6 | |
| bottom_part = np.ones((270, 2100, 3), dtype=np.uint8) * 255 | |
| elif size==34: | |
| img_per_row = 7 | |
| bottom_part = np.ones((268, 2100, 3), dtype=np.uint8) * 255 | |
| else: | |
| img_per_row = 10 | |
| bottom_part = np.ones((470, 2100, 3), dtype=np.uint8) * 255 | |
| img_2 = [] | |
| num_ = number | |
| if num_%img_per_row == 0: | |
| for i in range(int(num_/img_per_row)): | |
| img1 = horizontal_merge(img_per_row,w_output,size) | |
| img_2.append(img1) | |
| # print('fill') | |
| for i in range(img_per_row-int(num_/img_per_row)): | |
| img1_1 = horizontal_merge(0,w_output,size) | |
| img_2.append(img1_1) | |
| # print('blank') | |
| elif num_%img_per_row != 0: | |
| for i in range(int(num_//img_per_row)): | |
| img1 = horizontal_merge(img_per_row,w_output,size) | |
| img_2.append(img1) | |
| # print('fill') | |
| img1 = horizontal_merge(num_%img_per_row,w_output,size) | |
| img_2.append(img1) | |
| # print('some') | |
| for i in range(img_per_row-int(num_//img_per_row)-1): | |
| img1 = horizontal_merge(0,w_output,size) | |
| img_2.append(img1) | |
| # print('blank') | |
| merge_img_final = cv2.vconcat(img_2) | |
| # print(merge_img_final.shape) | |
| # bottom_part = np.ones((270, 2100, 3), dtype=np.uint8) * 255 | |
| merge_img_final = cv2.vconcat([merge_img_final,bottom_part]) | |
| # print(merge_img_final.shape) | |
| merge_img_final = add_border(merge_img_final, 52, 32, 30, 30, (255,255,255)) | |
| # print(merge_img_final.shape) | |
| # cv2.imwrite(output_path, merge_img_final) | |
| # print(merge_img_final.shape) | |
| 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 |