Update app.py

app.py

@@ -1,21 +1,72 @@
 import gradio as gr
 import torch
 from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq
-from PIL import Image
+from PIL import Image, ImageOps
 import numpy as np
 import os
 from huggingface_hub import snapshot_download
 import logging
+from pathlib import Path
+import tempfile
+import requests
+from io import BytesIO
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+class ImageHandler:
+    """Handle image processing and conversion"""
+    @staticmethod
+    def convert_to_rgb(image_path):
+        """Convert image to RGB format supporting multiple formats"""
+        try:
+            # If image is a URL, download it first
+            if isinstance(image_path, str) and (image_path.startswith('http://') or image_path.startswith('https://')):
+                response = requests.get(image_path)
+                image_data = BytesIO(response.content)
+                image = Image.open(image_data)
+            else:
+                image = Image.open(image_path)
+
+            # Convert RGBA to RGB if needed
+            if image.mode == 'RGBA':
+                background = Image.new('RGB', image.size, (255, 255, 255))
+                background.paste(image, mask=image.split()[3])
+                image = background
+            # Convert any other mode to RGB
+            elif image.mode != 'RGB':
+                image = image.convert('RGB')
+                
+            return image
+            
+        except Exception as e:
+            logger.error(f"Error converting image: {str(e)}")
+            raise ValueError(f"Não foi possível processar a imagem. Erro: {str(e)}")
+
+    @staticmethod
+    def process_image(image):
+        """Process image from various input types"""
+        try:
+            if isinstance(image, np.ndarray):
+                return Image.fromarray(image)
+            elif isinstance(image, Image.Image):
+                return image
+            elif isinstance(image, (str, Path)):
+                return ImageHandler.convert_to_rgb(image)
+            else:
+                raise ValueError("Formato de imagem não suportado")
+                
+        except Exception as e:
+            logger.error(f"Error processing image: {str(e)}")
+            raise ValueError(f"Erro no processamento da imagem: {str(e)}")
+
 class NutritionalAnalyzer:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.models = {}
         self.processors = {}
+        self.image_handler = ImageHandler()
         
     def initialize_model(self, model_name):
         """Initialize a specific model with error handling and caching"""
@@ -37,7 +88,7 @@ class NutritionalAnalyzer:
                 
                 config = model_configs.get(model_name)
                 if not config:
-                    raise ValueError(f"Unsupported model: {model_name}")
+                    raise ValueError(f"Modelo não suportado: {model_name}")
                 
                 # Ensure cache directory exists
                 os.makedirs(config["local_cache"], exist_ok=True)
@@ -72,19 +123,6 @@ class NutritionalAnalyzer:
             logger.error(f"Error initializing {model_name}: {str(e)}")
             return False
 
-    def prepare_image(self, image):
-        """Prepare image for model input"""
-        if isinstance(image, str):
-            image = Image.open(image)
-        elif isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
-        
-        # Ensure image is in RGB mode
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-            
-        return image
-
     def generate_nutritional_prompt(self, user_question):
         """Generate a comprehensive nutritional analysis prompt"""
         return f"""Como nutricionista especializado, analise esta refeição detalhadamente:
@@ -111,6 +149,9 @@ Por favor, forneça uma análise detalhada em português."""
     def analyze_image(self, image, question, model_choice):
         """Analyze image with nutritional focus"""
         try:
+            if image is None:
+                return "Por favor, faça upload de uma imagem para análise."
+
             # Convert model choice to internal name
             model_name = model_choice.lower().replace("-", "")
             
@@ -118,34 +159,48 @@ Por favor, forneça uma análise detalhada em português."""
             if not self.initialize_model(model_name):
                 return "Erro: Não foi possível inicializar o modelo. Por favor, tente novamente."
             
-            # Prepare image and prompt
-            processed_image = self.prepare_image(image)
+            # Process image with enhanced error handling
+            try:
+                processed_image = self.image_handler.process_image(image)
+            except ValueError as e:
+                return str(e)
+            except Exception as e:
+                return f"Erro no processamento da imagem: {str(e)}"
+
+            # Generate and process prompt
             nutritional_prompt = self.generate_nutritional_prompt(question)
             
             # Process input
-            inputs = self.processors[model_name](
-                images=processed_image,
-                text=nutritional_prompt,
-                return_tensors="pt"
-            ).to(self.device)
+            try:
+                inputs = self.processors[model_name](
+                    images=processed_image,
+                    text=nutritional_prompt,
+                    return_tensors="pt"
+                ).to(self.device)
+            except Exception as e:
+                return f"Erro no processamento do modelo: {str(e)}"
             
             # Generate response with enhanced parameters
-            with torch.no_grad():
-                outputs = self.models[model_name].generate(
-                    **inputs,
-                    max_new_tokens=300,
-                    num_beams=5,
-                    temperature=0.7,
-                    top_p=0.9,
-                    repetition_penalty=1.2,
-                    length_penalty=1.0
-                )
-            
-            # Decode and format response
-            response = self.processors[model_name].decode(outputs[0], skip_special_tokens=True)
-            formatted_response = self.format_response(response)
+            try:
+                with torch.no_grad():
+                    outputs = self.models[model_name].generate(
+                        **inputs,
+                        max_new_tokens=300,
+                        num_beams=5,
+                        temperature=0.7,
+                        top_p=0.9,
+                        repetition_penalty=1.2,
+                        length_penalty=1.0
+                    )
+                
+                # Decode and format response
+                response = self.processors[model_name].decode(outputs[0], skip_special_tokens=True)
+                formatted_response = self.format_response(response)
+                
+                return formatted_response
             
-            return formatted_response
+            except Exception as e:
+                return f"Erro na geração da análise: {str(e)}"
             
         except Exception as e:
             logger.error(f"Analysis error: {str(e)}")
@@ -234,6 +289,7 @@ def create_interface():
             2. Capture todos os elementos do prato
             3. Evite ângulos muito inclinados
             4. Seja específico em suas perguntas
+            5. Formatos de imagem suportados: JPG, PNG, WEBP
             """)
         
         analyze_btn.click(