Update app.py

app.py

@@ -30,19 +30,23 @@ def download_sam_model():
                 progress_bar.update(size)
     return model_path
 
-def show_mask(mask, image, random_color=True):
+def show_mask(mask, image):
     """Aplica uma máscara colorida sobre a imagem"""
-    if random_color:
-        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
-    else:
-        color = np.array([30/255, 144/255, 255/255, 0.6])
+    # Gerar cor aleatória para a máscara
+    color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    
     h, w = mask.shape[-2:]
     mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
     
-    img_rgba = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+    # Converter imagem para RGBA se necessário
+    if image.shape[-1] == 3:
+        img_rgba = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
+    else:
+        img_rgba = image
+    
     mask_image_uint8 = (mask_image * 255).astype(np.uint8)
     
-    # Mistura a máscara com a imagem original
+    # Misturar a máscara com a imagem original
     return cv2.addWeighted(img_rgba, 0.7, mask_image_uint8, 0.3, 0)
 
 def load_sam_model():
@@ -96,24 +100,25 @@ def segment_teeth(image, predictor):
     
     # Gerar pontos para segmentação
     input_points = generate_grid_points(image)
-    input_labels = np.ones(len(input_points))  # 1 para todos os pontos
+    input_labels = np.ones(len(input_points))
     
-    # Gerar máscaras
+    # Gerar máscaras usando os parâmetros corretos
     masks, scores, logits = predictor.predict(
         point_coords=input_points,
         point_labels=input_labels,
-        multimask_output=True,
-        num_multimask_outputs=3
+        multimask_output=True  # Isso permitirá múltiplas máscaras por ponto
     )
     
     # Criar visualização da segmentação
     final_image = image.copy()
-    for i, mask in enumerate(masks):
+    
+    # Aplicar cada máscara com uma cor diferente
+    for mask in masks:
         final_image = show_mask(mask, final_image)
-        
-        # Desenhar pontos de entrada
-        for point in input_points:
-            cv2.circle(final_image, (int(point[0]), int(point[1])), 3, (255, 0, 0), -1)
+    
+    # Desenhar os pontos de referência
+    for point in input_points:
+        cv2.circle(final_image, (int(point[0]), int(point[1])), 3, (255, 0, 0), -1)
     
     return final_image, masks
 
@@ -159,6 +164,9 @@ Observações:
         return result, analysis
         
     except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        print(error_trace)  # Isso ajudará na depuração
         return None, f"Erro durante o processamento: {str(e)}"
 
 # Create Gradio interface