Salvando alterações

Inferencia.py

@@ -1,5 +1,7 @@
 from transformers import pipeline
 import torch
+import json
+
 class Inferencia:
 
     def __init__(self):
@@ -28,4 +30,27 @@ class Inferencia:
 
         output = pipe(text=messages, max_new_tokens=200)
 
-        return output[0]["generated_text"][-1]["content"]
+        return output[0]["generated_text"][-1]["content"]
+
+    @staticmethod
+    def string_para_dicionario(string_json):
+        """
+        Transforma uma string JSON envolta em crases e com a palavra 'json' em um dicionário Python.
+
+        Args:
+            string_json (str): A string JSON a ser convertida, podendo estar envolta em crases e conter 'json'.
+
+        Returns:
+            dict: Um dicionário Python correspondente à string JSON.
+        """
+        try:
+            # Remove as crases (`) e a palavra 'json' se estiverem presentes
+            if string_json.startswith("```json") and string_json.endswith("```"):
+                string_json = string_json[7:-3].strip()  # Remove "```json" e "```"
+
+            # Converte a string JSON limpa em um dicionário Python
+            dicionario = json.loads(string_json)
+            return dicionario
+        except json.JSONDecodeError as e:
+            # Caso ocorra um erro na conversão, levanta uma exceção com detalhes
+            raise ValueError(f"Erro ao decodificar a string JSON: {e}")

PromptDocumento.py

@@ -0,0 +1,32 @@
+
+class PromptDocumento:
+
+    @staticmethod
+    def cnh():
+        return """
+        Extraia o nome, cpf, data de nascimento, número de registro, primeira habilitação,
+        validade da CNH, RG (documento de identidade), órgão emissor, estado,
+        categoria da habilitação, nacionalidade, filiação pai (se houver), filiação mãe (se houver),
+        local, data da emissão e observações (se houver).
+        Para os campos "se houver", caso não tenha nenhum valor, o JSON será preenchido em branco.
+        Retorne apenas o JSON abaixo preenchido com os valores encontrados:
+        
+        {
+            "nome": "",
+            "cpf": "",
+            "data_nascimento": "",
+            "numero_registro": "",
+            "primeiro_habilitacao": "",
+            "validade_cnh": "",
+            "rg": "",
+            "orgao_emissor": "",
+            "orgao_emissor_estado": "",
+            "categoria_habilitacao": "",
+            "nacionalidade": "",
+            "filiacao_pai": "",
+            "filiacao_mae": "",
+            "local": "",
+            "data_emissao": "",
+            "observacoes": "",
+        }         
+        """

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 from LogCustomizado import LogCustomizado as logc
 from Inferencia import Inferencia
 from PIL import Image
+from PromptDocumento import PromptDocumento
 
 st.set_page_config(page_title="CADIN - Cadastro Inteligente")
 
@@ -34,11 +35,21 @@ if arquivos_upload is not None:
             prompt_analisa_tipo_documento,
             imagem_convertida
         )
+
         st.write(dados_imagem)
+
         logc.gerar_log(mensagem_log=f"DADOS IMAGEM: {dados_imagem}")
 
-        # "Extraia o nome, cpf, data de nascimento e número de registro. Retorne apenas um JSON com esses dados."
+        dados_imagem_dict = Inferencia.string_para_dicionario(dados_imagem)
 
-        logc.gerar_log(mensagem_log=f"ARQUIVO PROCESSADO: {arquivo.name}")
+        if dados_imagem_dict["tipo_documento"] == "CNH":
+            dados_cnh = Inferencia.extrair_dados_imagem(
+                PromptDocumento.cnh(),
+                imagem_convertida
+            )
+            logc.gerar_log(mensagem_log=f"DADOS CNH: {dados_cnh}")
+            logc.gerar_log(mensagem_log=f"ARQUIVO PROCESSADO: {arquivo.name}")
 
     logc.gerar_log(mensagem_log=f"FIM DO PROCESSAMENTO")
+
+