Update app.py

app.py

@@ -5,29 +5,22 @@ from PIL import Image
 import numpy as np
 
 WOUND_TYPES = {
-    "pressure_ulcer": "Úlcera de Pressão",
-    "venous_ulcer": "Úlcera Venosa",
-    "diabetic_ulcer": "Úlcera Diabética",
-    "surgical_wound": "Ferida Cirúrgica",
-    "burn": "Queimadura",
-    "healthy_skin": "Pele Saudável"
+    "stage1": "Estágio 1 - Lesão Superficial",
+    "stage2": "Estágio 2 - Lesão Parcial",
+    "stage3": "Estágio 3 - Lesão Profunda",
+    "stage4": "Estágio 4 - Lesão Grave",
+    "unstageable": "Não Classificável",
+    "healthy": "Pele Saudável"
 }
 
-def load_models():
-    # Carregando modelos específicos para feridas
-    wound_classifier = pipeline(
+def load_model():
+    # Usando um modelo público do Hugging Face para classificação de imagens
+    classifier = pipeline(
         "image-classification",
-        model="stevhliu/wound-classification",
+        model="google/vit-base-patch16-224",
         device=0 if torch.cuda.is_available() else -1
     )
-    
-    tissue_classifier = pipeline(
-        "image-classification",
-        model="viktorcikojevic/wound-tissue-type",
-        device=0 if torch.cuda.is_available() else -1
-    )
-    
-    return wound_classifier, tissue_classifier
+    return classifier
 
 def preprocess_image(image):
     if isinstance(image, np.ndarray):
@@ -37,38 +30,29 @@ def preprocess_image(image):
 
 def classify_wound(image):
     if image is None:
-        return None, None
+        return None
     
-    wound_classifier, tissue_classifier = load_models()
+    classifier = load_model()
     processed_image = preprocess_image(image)
     
-    # Classificação do tipo de ferida
-    wound_results = wound_classifier(processed_image)
-    
-    # Classificação do tipo de tecido
-    tissue_results = tissue_classifier(processed_image)
+    # Classificação da imagem
+    results = classifier(processed_image)
     
     # Formatando resultados
-    wound_formatted = []
-    for result in wound_results:
-        label = WOUND_TYPES.get(result['label'], result['label'])
-        score = result['score']
-        wound_formatted.append((label, score))
-    
-    tissue_formatted = []
-    for result in tissue_results:
+    formatted_results = []
+    for result in results:
         label = result['label'].replace('_', ' ').title()
         score = result['score']
-        tissue_formatted.append((label, score))
+        formatted_results.append((label, score))
     
-    return wound_formatted, tissue_formatted
+    return formatted_results
 
 # Interface Gradio
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🏥 Classificador Especializado de Feridas
+    # 🏥 Classificador de Imagens Médicas
     
-    Sistema de classificação utilizando modelos especializados para análise de feridas e tecidos.
+    Sistema de classificação de imagens usando Vision Transformer (ViT).
     """)
     
     with gr.Row():
@@ -77,46 +61,34 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Upload da Imagem",
                 type="pil"
             )
-            submit_btn = gr.Button("Analisar Ferida", variant="primary", size="lg")
+            submit_btn = gr.Button("Analisar Imagem", variant="primary")
         
         with gr.Column():
-            wound_output = gr.Label(
-                label="Tipo de Ferida",
-                num_top_classes=3
-            )
-            tissue_output = gr.Label(
-                label="Tipo de Tecido",
+            output = gr.Label(
+                label="Classificação",
                 num_top_classes=3
             )
     
     with gr.Row():
-        with gr.Accordion("Informações Importantes", open=False):
+        with gr.Accordion("Informações", open=False):
             gr.Markdown("""
-            ### Tipos de Feridas Detectáveis:
-            - Úlcera de Pressão
-            - Úlcera Venosa
-            - Úlcera Diabética
-            - Ferida Cirúrgica
-            - Queimadura
-            
-            ### Tipos de Tecidos Analisados:
-            - Tecido de Granulação
-            - Tecido Necrótico
-            - Tecido Fibrinoso
-            - Tecido Epitelial
-            
             ### Recomendações para Melhores Resultados:
             1. Use imagens bem iluminadas
-            2. Capture a imagem em um ângulo perpendicular à ferida
+            2. Capture a imagem em um ângulo perpendicular
             3. Mantenha um fundo neutro e limpo
             4. Evite sombras ou reflexos excessivos
+            
+            ### Observações:
+            - Este é um modelo de classificação geral
+            - Os resultados são aproximações e não substituem avaliação médica
+            - Consulte sempre um profissional de saúde para diagnóstico
             """)
     
     # Configurando eventos
     submit_btn.click(
         fn=classify_wound,
         inputs=input_image,
-        outputs=[wound_output, tissue_output]
+        outputs=output
     )
 
 if __name__ == "__main__":