aiutils.py

import cv2
from ultralytics import YOLO
from typing import List
import pytesseract
from PIL import Image
import logging

MODEL_PATH = "photo.pt"

# Global model object
model = YOLO(MODEL_PATH)
print('Model is loaded')


def get_photo_annotations(image_path, box_coords):
    # Read the input image
    image = cv2.imread(image_path)
    coords = []

    # Get image dimensions
    height, width, _ = image.shape

    # Convert and draw bounding boxes on the image
    for box_coord in box_coords:
        x_min, y_min, x_max, y_max = box_coord
        x1, y1, x2, y2 = int(x_min * width), int(y_min * height), int(x_max * width), int(y_max * height)
        coords.append([height, width, x1, y1, x2, y2])

    return coords, image


def run_ocr(image):
	try:
		# Convert cv2 image to PIL format
		image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
		# Run OCR using pytesseract
		text = pytesseract.image_to_string(image_pil)
		return text
	except Exception as e:
		logging.error('Exception in run_ocr(): ' + str(e))
		return ""


def clean_ocr(txt, add_space = True):
	ret = ""
	digit_counter = 0
	space_counter = 0

	for c in txt:
		if c == ' ' and add_space:
			ret += ' '
			space_counter += 1

		elif c.isdigit():
			ret += c
			digit_counter += 1

		# Number of digits in aadhar
		if digit_counter == 12:
			return ret

	return ret

def convert_title_case(text):
	s = ''

	for c in text:
		if c.isalpha() or c == ' ':
			s += c

	split_space = s.split(' ')
	new_s = ''

	for word in split_space:
		new_s = new_s + word.capitalize() + ' '

	return new_s.strip()


def find_dob(text_lines):
	total_lines = len(text_lines)
	dob_line = ''
	dob_text = None
	
	if total_lines == 2 or total_lines == 3:
		dob_line = text_lines[1]
	elif total_lines == 4:
		dob_line = text_lines[2]
	else:
		dob_line = text_lines[4]

	space_split = dob_line.split(' ')[-1]

	len_space_split = len(space_split)
	i = len_space_split -1

	temp_dob = ''

	while i >= 0:
		if space_split[i].isdigit():
			temp_dob = space_split[i] + temp_dob
		if len(temp_dob) >= 8:
			break
		i -= 1

	if len(temp_dob) == 4:
		dob_text = temp_dob
	
	elif len(temp_dob) == 8:
		date = temp_dob[0:2]
		month = temp_dob[2:4]
		year = temp_dob[4:]

		if int(date) > 31:
			date = 'XX'
		
		if int(month) > 12:
			month = 'XX'
		
		if int(year[0] + year[1]) < 19 or int(year[0] + year[1]) > 20:
			year = 'XXXX'

		dob_text = date + '/' + month + '/' + year

	elif len(temp_dob) > 8:
		# Extract year only
		dob_text = temp_dob[-4] + temp_dob[-3] + temp_dob[-2] + temp_dob[-1]

	return dob_text


def find_dob2(text_lines):
	dob_line = None
	# year, dob, numbers, date, birth
	dob_line_index = 0
	dob_text = None

	try:
		for line in text_lines:
			if 'year' in line or \
				'dob' in line or \
				'date' in line or \
				'birth' in line:
				dob_line = line
				break
			else:
				dob_line_index += 1

		space_split = dob_line.split(' ')[-1]

		len_space_split = len(space_split)
		i = len_space_split -1

		temp_dob = ''

		while i >= 0:
			if space_split[i].isdigit():
				temp_dob = space_split[i] + temp_dob

			if len(temp_dob) >= 8:
				break

			i -= 1

		if len(temp_dob) == 4:
			dob_text = temp_dob
		
		elif len(temp_dob) == 8:
			dob_text = temp_dob[0:2] + '/' + temp_dob[2:4] + '/' + temp_dob[4:]

		elif len(temp_dob) > 8:
			# Extract year only
			dob_text = temp_dob[-4] + temp_dob[-3] + temp_dob[-2] + temp_dob[-1]
	
	except Exception as e:
		logging.error('Exception in find_dob2(): ' + str(e))

	return dob_text, dob_line_index


def find_gender(text):
	gender = None
	if 'female' in text or 'femal' in text:
		gender = 'female'
	elif 'male' in text:
		gender = 'male'
	elif 'transgender' in text or 'ransg' in text or 'nsgen' in text:
		gender = 'transgender'
	
	return gender


def find_name(text_lines):
	total_lines = len(text_lines)

	name_text = text_lines[0]

	if '0' in text_lines[0] or '1' in text_lines[1]:
		name_text = text_lines[1]

	if total_lines == 2 or total_lines == 3:
		name_text = text_lines[0]
	elif total_lines == 4:
		name_text = text_lines[1]
	elif total_lines >= 4:	# == 6
		name_text = text_lines[1]

	if len(name_text) > 0:
		name_text = name_text.strip()
		name_text = convert_title_case(name_text)

	return [name_text]


def clean_more_details(text):
	gender = None
	dob = None
	names = None

	text_lower = text.lower()

	# Gender
	try:
		gender = find_gender(text_lower)
	except Exception as e:
		logging.error('Exception in find_gender(): ' + str(e))

	text_lines_ocr = text_lower.split('\n')
	text_lines = []
	
	for line in text_lines_ocr:
		if len(line) > 0:
			text_lines.append(line)

	try:
		# dob, dob_line_index = find_dob2(text_lines)
		dob, dob_line_index = find_dob2(text_lines)
	except Exception as e:
		logging.error('Exception in find_dob2(): ' + str(e))

	try:
		if dob is None:
			dob = find_dob(text_lines)
	except Exception as e:
		logging.error('Exception in find_dob(): ' + str(e))


	# Name
	try:
		names = find_name(text_lines)
	except Exception as e:
		logging.error('Exception in find_name(): ' + str(e))

	return {
		'name': names[0],
		'dob': dob,
		'gender': gender
	}


def get_text_details(coords, image):
	coords = coords[0]
	h, w, x1, y1, x2, y2 = coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]

	# x -> left right
	# y -> top bottom

	x2_factor = (x2-x1) * 2.86
	y1_factor = (y2-y1) * 0.152
	y2_factor = (y2-y1) * 0.168

	# print('x2_factor: ' + str(x2_factor) + ' y1_factor: ' + str(y1_factor) + ' y2_factor: ' + str(y2_factor))

	text_x1 = x2 - 5
	text_y1 = int(y1 - y1_factor)

	# correct text_x2 = max(int(text_x1 + x2_factor), photo_width - (photo_width * 0.18))
	text_x2 = min(int(text_x1 + x2_factor), int(w - (w * 0.08)))
	text_y2 = int(y2 - y2_factor)

	# print(str(text_x1) + ',' + str(text_y1) + ' ' + str(text_x2) + ',' + str(text_y2))

	# color = (0, 0, 255)
	# thickness = 2
	# cv2.rectangle(image, (text_x1, text_y1), (text_x2, text_y2), color, thickness)

	cropped_image = image[text_y1:text_y2, text_x1:text_x2]
	ocr_text = run_ocr(cropped_image)

	# print('ocr: ' + ocr_text)

	more_details = clean_more_details(ocr_text)
	return more_details



def get_aadhar_numbers(coords, image):
    coords = coords[0]
    h, w, x1, y1, x2, y2 = coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]

    photo_width = x2 - x1
    photo_height = y2 - y1

	# x -> left right
	# y -> top bottom

    factors = [
		# Case1: Image has QR, no Aadhar logo
		{
			'x1': photo_width / 11,
			'y1': photo_height / 14,
			'x2': photo_width * 2.2,
			'y2': photo_height / 4.5,
			'hwLow': 0,
			'hwHigh': 0.632
		},

		# Case2: No QR, Aadhar logo
		{
			'x1': photo_width / 10,
			'y1': photo_height / 2.4,
			'x2': photo_width * 2.2,
			'y2': photo_height / 4,
			'hwLow': 0.632,
			'hwhigh': 1
		},

		# Case3: Variation of Case2
		{
			'x1': photo_width / 10,
			'y1': photo_height / 2,
			'x2': photo_width * 2.2,
			'y2': photo_height / 3.1,
			'hwLow': 0.632,
			'hwhigh': 1
		},
	]

    red_boxes = []
    cleaned_list = []

    passed = False
    approach = 2

    if approach == 2:
        factor_index = 0
        factor = factors[0]

        i_x1 = int(x2 + factor['x1'] - 20)
        i_y1 = int(y2 + factor['y1'])
        i_x2 = int(i_x1 + factor['x2']) 
        i_y2 = int(i_y1 + factor['y2'])

        runs = 17
        while runs > 0:
            runs -= 1
            cropped_image = image[i_y1:i_y2, i_x1:i_x2]

            ocr_text = run_ocr(cropped_image)
            cleaned_txt = clean_ocr(ocr_text, False)

            if len(cleaned_txt.replace(' ', '')) == 12:
                  cleaned_list.append(cleaned_txt)

            i_y1 = i_y1 + 10
            i_y2 = i_y2 + 10

    return cleaned_list


def extract_results(results):
    responses_full = []

    for result in results:
        response = dict()
        response['path'] = result.path
        response['box_coords'] = result.boxes.xyxyn.cpu().numpy().tolist()
        responses_full.append(response)

    return responses_full


def get_aadhar_number(file_path):
    resp = {
        'detectPhoto': False,
        'aadharNumber': None,
        'dob': None,
        'gender': None,
        'name': None
    }

    try:
        model_results = model(file_path, verbose=False)
        photo_loc_data = extract_results(model_results)

        coords, image = get_photo_annotations(photo_loc_data[0]['path'], photo_loc_data[0]['box_coords'])
        try:
            cleaned_numbers = get_aadhar_numbers(coords, image)

            has_photo = False

            if len(coords) > 0:
                has_photo = True

            resp['detectPhoto'] = has_photo
            resp['aadharNumber'] = cleaned_numbers[0]
        except Exception as e:
            logging.error('Exception: ' + str(e))

        more_details = get_text_details(coords, image)
        resp['name'] = more_details['name']
        resp['dob'] = more_details['dob']
        resp['gender'] = more_details['gender']
       
    except Exception as e:
        logging.error('Exception in get_aadhar_number() : ' + str(e))

    return resp