Update app.py

app.py

@@ -16,7 +16,6 @@ CLASS_NAMES = {
     3: "packaging",
 }
 
-# Caminho para o modelo
 MODEL_PATH = "segmentation_model.h5"
 model = tf.keras.models.load_model(MODEL_PATH)
 
@@ -40,19 +39,22 @@ def predict_image(input_image):
     original_size = original_img_np.shape[:2]
     mask_final = cv2.resize(mask_predicted.astype(np.uint8), (original_size[1], original_size[0]), interpolation=cv2.INTER_NEAREST)
     confidences_final = cv2.resize(confidences, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR)
-
+    
+    # Processa a imagem para desenhar as caixas
     final_img = original_img_cv2.copy()
-    confidence_threshold = 0.7  # Limiar de 80%
+    confidence_threshold = 0.8  # Limiar de 80%
 
     for class_id in np.unique(mask_final):
         if class_id == 0:
             continue
         
         class_name = CLASS_NAMES.get(class_id, f"Classe {class_id}")
-        binary_mask = (mask_final == class_id).astype(np.uint8)
+        
+        binary_mask = (mask_final == class_id)
         binary_mask = binary_fill_holes(binary_mask)
+        binary_mask_uint8 = binary_mask.astype(np.uint8)
         
-        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        contours, _ = cv2.findContours(binary_mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
         for contour in contours:
             x_min, y_min, w, h = cv2.boundingRect(contour)
@@ -70,17 +72,19 @@ def predict_image(input_image):
                 label_text = f"{class_name}: {avg_confidence:.2f}%"
                 color_tuple = CLASS_COLORS.get(class_id, (255, 255, 255))
                 
+                # Desenha a caixa
                 cv2.rectangle(final_img, (x_min, y_min), (x_max, y_max), color_tuple, 2)
                 
+                # Desenha o fundo do texto
                 (text_width, text_height), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
                 cv2.rectangle(final_img, (x_min, y_min - text_height - 10), (x_min + text_width, y_min), color_tuple, -1)
                 cv2.putText(final_img, label_text, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
     
-
+    # Converte de volta para PIL e retorna
     final_img_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
     return Image.fromarray(final_img_rgb)
 
-
+# Define a interface Gradio
 gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),