Spaces:
Sleeping
Sleeping
import os | |
from google.cloud import vision | |
import re | |
import torch | |
import torchvision | |
import numpy as np | |
from PIL import Image | |
import albumentations as A | |
from albumentations.pytorch import ToTensorV2 | |
import tempfile | |
import json | |
def getcredentials(): | |
secret_key_credential = os.getenv("cloud_vision") | |
with tempfile.NamedTemporaryFile(mode='w+', delete= False, suffix=".json") as temp_file: | |
temp_file.write(secret_key_credential) | |
tempfile_name = temp_file.name | |
return tempfile_name | |
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = getcredentials() | |
## | |
def info_new_cni(donnees): | |
## | |
informations = {} | |
# Utilisation d'expressions régulières pour extraire les informations spécifiques | |
numero_carte = re.search(r'n° (C\d+)', ' '.join(donnees)) | |
#prenom_nom = re.search(r'Prénom\(s\)\s+(.*?)\s+Nom\s+(.*?)\s+Signature', ' '.join(donnees)) | |
nom = re.search(r'Nom\s+(.*?)\s', ' '.join(donnees)) | |
prenom = re.search(r'Prénom\(s\)\s+(.*?)\s+Nom\s+(.*?)', ' '.join(donnees)) | |
date_naissance = re.search(r'Date de Naissance\s+(.*?)+(\d{2}/\d{2}/\d{4})', ' '.join(donnees)) | |
lieu_naissance = re.search(r'Lieu de Naissance\s+(.*?)\s', ' '.join(donnees)) | |
taille = re.search(r'Sexe Taille\s+(.*?)+(\d+,\d+)', ' '.join(donnees)) | |
nationalite = re.search(r'Nationalité\s+(.*?)\s+\d+', ' '.join(donnees)) | |
date_expiration = re.search(r'Date d\'expiration\s+(\d+/\d+/\d+)', ' '.join(donnees)) | |
sexe = re.search(r'Date de Naissance\s+(.*?)+(\d{2}/\d{2}/\d{4})+(.*)', ' '.join(donnees)) | |
# Stockage des informations extraites dans un dictionnaire | |
if numero_carte: | |
informations['Numéro de carte'] = numero_carte.group(1) | |
if nom : | |
informations['Nom'] = nom.group(1) | |
if prenom: | |
informations['Prénom'] = prenom.group(1) | |
if date_naissance: | |
informations['Date de Naissance'] = date_naissance.group(2) | |
if lieu_naissance: | |
informations['Lieu de Naissance'] = lieu_naissance.group(1) | |
if taille: | |
informations['Taille'] = taille.group(2) | |
if nationalite: | |
informations['Nationalité'] = nationalite.group(1) | |
if date_expiration: | |
informations['Date d\'expiration'] = date_expiration.group(1) | |
if sexe : | |
informations['sexe'] = sexe.group(3)[:2] | |
return informations | |
## | |
def info_ancien_cni(infos): | |
""" Extract information in row data of ocr""" | |
informations = {} | |
immatriculation_patern = r'Immatriculation:\s+(C \d{4} \d{4} \d{2})' | |
immatriculation = re.search(immatriculation_patern, ''.join(infos)) | |
nom = infos[4] | |
prenom_pattern = r'Nom\n(.*?)\n' | |
prenom = re.search(prenom_pattern, '\n'.join(infos)) | |
sexe_pattern = r'Prénoms\n(.*?)\n' | |
sexe = re.search(sexe_pattern, '\n'.join(infos)) | |
taille_pattern = r'Sexe\n(.*?)\n' | |
taille = re.search(taille_pattern, '\n'.join(infos)) | |
date_naiss_pattern = r'Taille\s+(.*?)+(\d+/\d+/\d+)' # r'Taille (m)\n(.*?)\n' | |
date_naissance = re.search(date_naiss_pattern, ' '.join(infos)) | |
lieu_pattern = r'Date de Naissance\n(.*?)\n' | |
lieu_naissance = re.search(lieu_pattern, '\n'.join(infos)) | |
valide_pattern = r'Valide jusqu\'au+(.*?)+(\d+/\d+/\d+)' | |
validite = re.search(valide_pattern, ' '.join(infos)) | |
# Stockage des informations extraites dans un dictionnaire | |
if immatriculation: | |
informations['Immatriculation'] = immatriculation.group(1) | |
if nom : | |
informations['Nom'] = infos[4] | |
if prenom: | |
informations['Prénom'] = prenom.group(1) | |
if date_naissance: | |
informations['Date de Naissance'] = date_naissance.group(2) | |
if lieu_naissance: | |
informations['Lieu de Naissance'] = lieu_naissance.group(1) | |
if taille: | |
informations['Taille'] = taille.group(1) | |
if validite: | |
informations['Date d\'expiration'] = validite.group(2) | |
if sexe : | |
informations['sexe'] = sexe.group(1) | |
return informations | |
## | |
def filtrer_elements(liste): | |
elements_filtres = [] | |
for element in liste: | |
if element not in ['\r',"RÉPUBLIQUE DE CÔTE D'IVOIRE", "MINISTÈRE DES TRANSPORTS", "PERMIS DE CONDUIRE"]: | |
elements_filtres.append(element) | |
return elements_filtres | |
def permis_de_conduite(donnees): | |
""" Extraire les information de permis de conduire""" | |
informations = {} | |
infos = filtrer_elements(donnees) | |
nom_pattern = r'Nom\n(.*?)\n' | |
nom = re.search(nom_pattern, '\n'.join(infos)) | |
prenom_pattern = r'Prénoms\n(.*?)\n' | |
prenom = re.search(prenom_pattern, '\n'.join(infos)) | |
date_lieu_naissance_patern = r'Date et lieu de naissance\n(.*?)\n' | |
date_lieu_naissance = re.search(date_lieu_naissance_patern, '\n'.join(infos)) | |
date_lieu_delivrance_patern = r'Date et lieu de délivrance\n(.*?)\n' | |
date_lieu_delivrance = re.search(date_lieu_delivrance_patern, '\n'.join(infos)) | |
numero_pattern = r'Numéro du permis de conduire\n(.*?)\n' | |
numero = re.search(numero_pattern, '\n'.join(infos)) | |
restriction_pattern = r'Restriction\(s\)\s+(.*?)+(.*)' | |
restriction = re.search(restriction_pattern, ' '.join(infos)) | |
# Stockage des informations extraites dans un dictionnaire | |
if nom: | |
informations['Nom'] = nom.group(1) | |
if prenom : | |
informations['Prenoms'] = prenom.group(1) | |
if date_lieu_naissance : | |
informations['Date_et_lieu_de_naissance'] = date_lieu_naissance.group(1) | |
if date_lieu_naissance : | |
informations['Date_et_lieu_de_délivrance'] = date_lieu_delivrance.group(1) | |
informations['Categorie'] = infos[0] | |
if numero: | |
informations['Numéro_du_permis_de_conduire'] = numero.group(1) | |
if restriction: | |
informations['Restriction(s)'] = restriction.group(2) | |
return informations | |
# Fonction pour extraire les informations individuelles | |
def extraire_informations_carte(path, type_de_piece=1): | |
""" Detect text in identity card""" | |
client = vision.ImageAnnotatorClient() | |
with open(path,'rb') as image_file: | |
content = image_file.read() | |
image = vision.Image(content = content) | |
# for non dense text | |
#response = client.text_detection(image=image) | |
#for dense text | |
response = client.document_text_detection(image = image) | |
texts = response.text_annotations | |
ocr_texts = [] | |
for text in texts: | |
ocr_texts.append(f"\r\n{text.description}") | |
if response.error.message : | |
raise Exception("{}\n For more informations check : https://cloud.google.com/apis/design/errors".format(response.error.message)) | |
donnees = ocr_texts[0].split('\n') | |
if type_de_piece ==1: | |
return info_new_cni(donnees) | |
elif type_de_piece == 2: | |
return info_ancien_cni(donnees) | |
elif type_de_piece == 3: | |
return permis_de_conduite(donnees) | |
else : | |
return "Le traitement de ce type de document n'est pas encore pris en charge" | |
def load_checkpoint(path): | |
print('--> Loading checkpoint') | |
return torch.load(path,map_location=torch.device('cpu')) | |
def make_prediction(image_path): | |
# define the using of GPU or CPU et background training | |
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
## load model | |
model = load_checkpoint("data/model.pth") | |
## transformation | |
test_transforms = A.Compose([ | |
A.Resize(height=224, width=224, always_apply=True), | |
A.Normalize(always_apply=True), | |
ToTensorV2(always_apply=True),]) | |
## read the image | |
image = np.array(Image.open(image_path).convert('RGB')) | |
transformed = test_transforms(image= image) | |
image_transformed = transformed["image"] | |
image_transformed = image_transformed.unsqueeze(0) | |
image_transformed = image_transformed.to(device) | |
model.eval() | |
with torch.set_grad_enabled(False): | |
output = model(image_transformed) | |
# Post-process predictions | |
probabilities = torch.nn.functional.softmax(output[0], dim=0) | |
predicted_class = torch.argmax(probabilities).item() | |
proba = float(max(probabilities)) | |
return proba, predicted_class |