detection.py

from ultralytics import YOLO
import cv2
from typing import List
import pytesseract
from PIL import 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:
		print('Exception: ' + str(e))
		return ""


def clean_ocr_number(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 draw_annotations_on_image(image_path, box_coords, output_path, draw=True):
	# 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)

        # Draw bounding box on the image
		color = (0, 255, 0)  # Green color
		thickness = 2
		if draw:
			cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)
		coords.append([height, width, x1, y1, x2, y2])

    # Show image
    # cv2.imshow('out', image)
    # cv2.waitKey()
    
    # Save the result
	cv2.imwrite(output_path, image)
	return coords, image



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_smallest_box(coords):
	min_height = 100000
	min_index = 0
	count = len(coords)

	for i in range(0, count):
		coord = coords[i]
		h, w, x1, y1, x2, y2 = coord[0], coord[1], coord[2], coord[3], coord[4], coord[5]

		if (y2-y1) < min_height:
			min_height = y2- y1
			min_index = i

	return coords[min_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 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()


'''
		Expected (num of lines) to text mappings
		2: eng name, dob, gender
		3: eng name, dob, gender
		4: hindi name, eng name, dob, gender
		6: hindi name, eng name, father hindi name, father eng name, dob, 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 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]
		print('space_split: ' + space_split)

		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

		print('temp_dob: ' + temp_dob)
		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:
		print('Exception in find_dob2(): ' + str(e))

	return dob_text, dob_line_index


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]
	print('space_split: ' + space_split)

	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

	print('temp_dob: ' + temp_dob)
	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 clean_more_details(text):
	gender = None
	dob = None
	names = None
	dob_line_index = 4

	text_lower = text.lower()

	# Gender
	try:
		gender = find_gender(text_lower)
	except Exception as e:
		print('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)
			print('line: '  + line)

	print('\n Count lines: ' + str(len(text_lines)))

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

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


	# Name
	try:
		names = find_name(text_lines)
	except Exception as e:
		print('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]

	photo_width = x2 - x1
	photo_height = y2 - y1

	hw_ratio = h / w
	wh_ratio = w / h

	# 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)
	print('More details: ')
	print(more_details)

	# Show image
	cv2.imshow('out', image)
	cv2.waitKey(0)

	return more_details




def draw_number_box(coords, image, correct_num):
	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

	hw_ratio = h / w
	wh_ratio = w / h

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

	# print('ph: ' + str(photo_height))
	# print('pw: ' + str(photo_width))
	# print('h-w ratio: ' + str(hw_ratio))
	# print('w-h ratio: ' + str(wh_ratio))

	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 = []

	passed = False
	approach = 2

	if approach == 1:
		factor_index = 0
		for factor in factors:
			print('Running for ' + str(factor_index))

			n_x1 = int(x2 + factor['x1'])
			n_y1 = int(y2 + factor['y1'])
			n_x2 = int(n_x1 + factor['x2'])
			n_y2 = int(n_y1 + factor['y2'])

			print('new: ' + str(n_x1) + ', ' + str(n_y1) + ', ' + str(n_x2) + ', ' + str(n_y2))

			color = (0, 0, 255)
			thickness = 2
			cv2.rectangle(image, (n_x1, n_y1), (n_x2, n_y2), color, thickness)

			cropped_image = image[n_y1:n_y2, n_x1:n_x2]

			ocr_text = run_ocr(cropped_image)
			cleaned_txt = clean_ocr_number(ocr_text)
			factor_text = "f_" + str(factor_index) + " raw: " + ocr_text + " clean: " + cleaned_txt

			# print('cleaned: ' + cleaned_txt)

			if cleaned_txt.replace(' ', '') == correct_num:
				passed = True

			font = cv2.FONT_HERSHEY_SIMPLEX
			font_scale = 0.5
			font_thickness = 1
			text_position = (n_x1, n_y1 - 5)  # Adjust the position to be just above the bounding box
			cv2.putText(image, factor_text, text_position, font, font_scale, color, font_thickness, cv2.LINE_AA)

			red_boxes.append([n_x1, n_y1, n_x1, n_y2])

			factor_index += 1

	elif approach == 2:
		factor_index = 0
		factor = factors[0]
		print('Running for ' + str(factor_index))

		i_x1 = int(x2 + factor['x1'])
		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
		
			# color = (0, 0, 255)
			# thickness = 2
			# cv2.rectangle(image, (i_x1, i_y1), (i_x2, i_y2), color, thickness)

			cropped_image = image[i_y1:i_y2, i_x1:i_x2]

			ocr_text = run_ocr(cropped_image)
			cleaned_txt = clean_ocr_number(ocr_text)
			factor_text = "f_" + str(factor_index) + " raw: " + ocr_text + " clean: " + cleaned_txt

			if len(cleaned_txt) > 0:
				print('cleaned: ' + cleaned_txt)

			if cleaned_txt.replace(' ', '') == correct_num:
				passed = True

			# font = cv2.FONT_HERSHEY_SIMPLEX
			# font_scale = 0.5
			# font_thickness = 1
			# text_position = (i_x1, i_y1 - 5)  # Adjust the position to be just above the bounding box
			# cv2.putText(image, factor_text, text_position, font, font_scale, color, font_thickness, cv2.LINE_AA)

			red_boxes.append([i_x1, i_y1, i_x1, i_y2])

			i_y1 = i_y1 + 10
			i_y2 = i_y2 + 10


	color = (255, 0, 0)
	cv2.putText(image, 'passed: ' + str(passed), (x2+50, y2-50), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)

	# Show image
	cv2.imshow('out', image)
	cv2.waitKey(0)

	return red_boxes


def get_address_details(coords, image):
	text = run_ocr(image)

	text_lines_ocr = text.split('\n')
	text_lines = []
	address_line = None

	for line in text_lines_ocr:
		if len(line) > 0:
			text_lines.append(line)

			if 'address' in line:
				address_line = line

			print('line: ' + line)
	
	print('\n address_line: ' + str(address_line))

	cv2.imshow('image', image)
	cv2.waitKey(0)



prediction_images = [
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/43.jpeg", "962527391229"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/66.jpeg", "373051268072"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/67.jpeg", "357936678550"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/68.jpeg", "357936678550"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/69.jpeg", "826049424804"],
	
	# ["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/70.jpeg", "957290915027"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/71.jpeg", "274454695540"],
	
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/1.jpeg", "272763753054"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/3.jpeg", "878823786726"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/4.jpeg", "853875354900"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/5.jpeg", "500075504173"],
	
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/6.jpeg", "387763345722"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/7.jpeg", "748224502548"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/8.jpeg", "958080582495"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/9.jpg", "253930176895"],
	
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/30.jpg", "436785640898"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/31.jpg", "222274457889"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/32.jpg", "494326098016"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/36.jpg", "274454695540"],
	
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/39.jpg", "662686835920"],
	["/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/img/40.jpeg", "700597966091"]
]

test_images = [
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/45.png', '439410712322'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/46.png', '574278330385'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/47.png', '547386153726'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/49.png', '352022273140'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/50.png', '211369809159'],

	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/51.png', '402109659501'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/53.png', '310454519517'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/54.png', '373051268072'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/55.png', '636354212716'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/57.png', '568968644949'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/58.png', '427016509926'],
	['/Users/anr/Documents/myFiles/dev/python/aadhar_flask/aadhar_data/test/60.png', '442217732001'],
]

# all_good_images = test_images
all_good_images = prediction_images

model = YOLO("/Users/anr/Documents/myFiles/dev/python/aadhar_flask/photo.pt")


for record in all_good_images:
	try:
		image_path = record[0]
		correct_num = record[1]

		result = model(image_path, verbose=False)

		parsed = extract_results(result)
		
		coords, image = draw_annotations_on_image(parsed[0]['path'], parsed[0]['box_coords'], "out.jpg", False)
		# print('\n coords: ' + str(coords))

		coords = [get_smallest_box(coords)]
		more_details = get_text_details(coords, image)
		# address_details = get_address_details(coords, image)
		number_coords = draw_number_box(coords, image, correct_num)
	
	except Exception as e:
		print('Exception: ' + str(e) + ' for image: ' + str(image_path))