update utils

utils.py

@@ -1,21 +1,14 @@
 import os
 from google.cloud import vision
 import re
-import json
-import tempfile
+import torch
+import torchvision
+import numpy as np
+from PIL import Image
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
 ##
-
-def get_credentials():
-    creds_json_str = os.getenv("cloud_vision")
-
-    #create temporale file
-    with tempfile.NamedTemporaryFile(mode="w+",delete=False, suffix=".json") as temp:
-        temp.write(creds_json_str) #write the content in json format
-        temp_filename = temp.name
-        
-    return temp_filename
-    
-os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = get_credentials()
+os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'data/ocr_vision_token.json'
 
 ##
 def info_new_cni(donnees):
@@ -116,14 +109,39 @@ def permis_de_conduite(donnees):
   """ Extraire les information de permis de conduire"""
 
   informations = {}
-  tab = filtrer_elements(donnees)
-  informations['Nom'] = tab[2]
-  informations['Prenoms'] = tab[4]
-  informations['Date_et_lieu_de_naissance'] = tab[6]
-  informations['Date_et_lieu_de_délivrance'] = tab[8]
-  informations['Categorie'] = tab[0]
-  informations['Numéro_du_permis_de_conduire'] = tab[10]
-  informations['Restriction(s)'] = tab[12:] if len(tab) > 11 else ''
+
+  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
 
@@ -161,4 +179,39 @@ def extraire_informations_carte(path, type_de_piece=1):
     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"
+      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