Update app.py

app.py

@@ -1,72 +1,73 @@
 import gradio as gr
 import torch
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForCausalLM
 import pandas as pd
-from PIL import Image
 import numpy as np
+from PIL import Image
 
-# Inicializa o modelo BLIP2
-device = "cuda" if torch.cuda.is_available() else "cpu"
-processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model = Blip2ForConditionalGeneration.from_pretrained(
-    "Salesforce/blip2-opt-2.7b",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    device_map="auto"
-)
+def get_model():
+    """Inicializa o modelo uma única vez"""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model_name = "microsoft/git-base-coco"
+    
+    processor = AutoProcessor.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = model.to(device)
+    
+    return processor, model, device
 
-# Base de dados nutricional
+# Base de dados nutricional simplificada
 NUTRITION_DB = {
     "arroz": {"calorias": 130, "proteinas": 2.7, "carboidratos": 28, "gorduras": 0.3},
     "feijão": {"calorias": 77, "proteinas": 5.2, "carboidratos": 13.6, "gorduras": 0.5},
     "frango": {"calorias": 165, "proteinas": 31, "carboidratos": 0, "gorduras": 3.6},
     "salada": {"calorias": 15, "proteinas": 1.4, "carboidratos": 2.9, "gorduras": 0.2},
-    "batata": {"calorias": 93, "proteinas": 2.5, "carboidratos": 21, "gorduras": 0.1},
-    "carne": {"calorias": 250, "proteinas": 26, "carboidratos": 0, "gorduras": 17},
-    "peixe": {"calorias": 206, "proteinas": 22, "carboidratos": 0, "gorduras": 12},
-    "macarrão": {"calorias": 158, "proteinas": 5.8, "carboidratos": 31, "gorduras": 1.2},
-    "ovo": {"calorias": 155, "proteinas": 13, "carboidratos": 1.1, "gorduras": 11},
 }
 
-def analyze_image(image, progress=gr.Progress()):
-    """Analisa a imagem usando BLIP2 e retorna descrição dos alimentos"""
+def process_image(image, progress=gr.Progress()):
+    """Processa a imagem e retorna a descrição"""
     try:
-        progress(0.2, desc="Processando imagem...")
+        progress(0.3, desc="Carregando modelo...")
+        processor, model, device = get_model()
         
-        # Garante que a imagem está no formato correto
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
+        progress(0.5, desc="Processando imagem...")
+        if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
-            
-        # Gera prompt específico para identificação de alimentos
-        prompt = "Identifique e liste todos os alimentos visíveis nesta imagem de refeição. Forneça uma lista detalhada."
-        
-        progress(0.4, desc="Analisando alimentos...")
-        inputs = processor(image, text=prompt, return_tensors="pt").to(device)
-        
-        progress(0.6, desc="Gerando descrição...")
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=100,
-                num_beams=5,
-                temperature=1.0,
-                top_p=0.9,
-            )
-            
-        food_description = processor.decode(outputs[0], skip_special_tokens=True)
-        progress(0.8, desc="Calculando nutrientes...")
         
-        # Identifica alimentos conhecidos na descrição
+        # Processa a imagem
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        
+        progress(0.7, desc="Gerando descrição...")
+        # Gera a descrição
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=50,
+            num_beams=1,
+            temperature=1.0,
+        )
+        
+        # Decodifica a saída
+        description = processor.decode(outputs[0], skip_special_tokens=True)
+        
+        progress(1.0, desc="Concluído!")
+        return description.strip()
+    
+    except Exception as e:
+        raise gr.Error(f"Erro no processamento da imagem: {str(e)}")
+
+def analyze_foods(description):
+    """Analisa a descrição e retorna informações nutricionais"""
+    try:
+        # Identifica alimentos da base de dados na descrição
         found_foods = []
         for food in NUTRITION_DB.keys():
-            if food in food_description.lower():
+            if food in description.lower():
                 found_foods.append(food)
         
         if not found_foods:
-            return None, None, None, "Não foi possível identificar alimentos conhecidos na imagem."
+            return "Nenhum alimento conhecido identificado.", None, None
         
-        # Calcula nutrientes totais
+        # Calcula nutrientes
         total_nutrients = {
             "calorias": 0,
             "proteinas": 0,
@@ -74,13 +75,12 @@ def analyze_image(image, progress=gr.Progress()):
             "gorduras": 0
         }
         
-        # Soma nutrientes de cada alimento encontrado
         for food in found_foods:
             for nutrient, value in NUTRITION_DB[food].items():
                 total_nutrients[nutrient] += value
         
         # Prepara dados para visualização
-        nutrients_table = [
+        table_data = [
             ["Calorias", f"{total_nutrients['calorias']:.1f} kcal"],
             ["Proteínas", f"{total_nutrients['proteinas']:.1f}g"],
             ["Carboidratos", f"{total_nutrients['carboidratos']:.1f}g"],
@@ -97,71 +97,79 @@ def analyze_image(image, progress=gr.Progress()):
             ]
         })
         
-        # Gera análise
         analysis = f"""### Alimentos Identificados:
-{food_description}
+• {', '.join(found_foods)}
 
-### Alimentos na Base de Dados:
-{', '.join(found_foods)}
+### Descrição do Modelo:
+{description}
 
 ### Análise Nutricional:
-- Valor Calórico: {total_nutrients['calorias']:.1f} kcal
-- Proporção de Macronutrientes:
-  • Proteínas: {total_nutrients['proteinas']:.1f}g
-  • Carboidratos: {total_nutrients['carboidratos']:.1f}g
-  • Gorduras: {total_nutrients['gorduras']:.1f}g
+• Calorias Totais: {total_nutrients['calorias']:.1f} kcal
+• Proteínas: {total_nutrients['proteinas']:.1f}g
+• Carboidratos: {total_nutrients['carboidratos']:.1f}g
+• Gorduras: {total_nutrients['gorduras']:.1f}g
 """
         
-        progress(1.0, desc="Concluído!")
-        return analysis, nutrients_table, plot_data, None
+        return analysis, table_data, plot_data
         
     except Exception as e:
-        return None, None, None, f"Erro na análise: {str(e)}"
+        raise gr.Error(f"Erro na análise: {str(e)}")
+
+def analyze_image(image):
+    """Função principal que coordena o processo de análise"""
+    try:
+        # Processa a imagem
+        description = process_image(image)
+        
+        # Analisa os alimentos
+        analysis, table_data, plot_data = analyze_foods(description)
+        
+        return analysis, table_data, plot_data
+        
+    except Exception as e:
+        return str(e), None, None
 
 # Interface Gradio
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
-    gr.Markdown("# 🍽️ Análise Nutricional com IA")
+    gr.Markdown("""
+    # 🍽️ Análise Nutricional com IA
+    Faça upload de uma foto do seu prato para análise nutricional.
+    """)
     
     with gr.Row():
-        with gr.Column(scale=1):
-            # Input
+        # Coluna de Input
+        with gr.Column():
             image_input = gr.Image(
                 type="pil",
                 label="Foto do Prato",
                 sources=["upload", "webcam"]
             )
-            analyze_btn = gr.Button("Analisar", variant="primary")
+            analyze_btn = gr.Button("📊 Analisar", variant="primary")
+        
+        # Coluna de Output
+        with gr.Column():
+            # Análise textual
+            output_text = gr.Markdown(label="Análise")
             
-        with gr.Column(scale=2):
-            # Output
-            error_output = gr.Markdown(visible=False)
+            # Tabela nutricional
+            output_table = gr.Dataframe(
+                headers=["Nutriente", "Quantidade"],
+                label="Informação Nutricional"
+            )
             
-            with gr.Tabs():
-                with gr.Tab("Análise"):
-                    analysis_output = gr.Markdown()
-                
-                with gr.Tab("Nutrientes"):
-                    nutrients_table = gr.Dataframe(
-                        headers=["Nutriente", "Quantidade"],
-                        label="Informação Nutricional"
-                    )
-                    nutrients_plot = gr.BarPlot(
-                        x="Nutriente",
-                        y="Quantidade",
-                        title="Macronutrientes (g)",
-                        height=300
-                    )
+            # Gráfico
+            output_plot = gr.BarPlot(
+                x="Nutriente",
+                y="Quantidade",
+                title="Macronutrientes (g)",
+                height=300
+            )
     
     # Eventos
     analyze_btn.click(
         fn=analyze_image,
         inputs=[image_input],
-        outputs=[
-            analysis_output,
-            nutrients_table,
-            nutrients_plot,
-            error_output
-        ]
+        outputs=[output_text, output_table, output_plot]
     )
 
 if __name__ == "__main__":