Update MeAI_Maincode/Polyb.py

MeAI_Maincode/Polyb.py

@@ -7,11 +7,11 @@ from lang_sam import LangSAM
 model = LangSAM()
 
 text_prompt = """
-A polyp is an anomalous oval-shaped small bump-like structure, a relatively small growth or
-mass that develops on the inner lining of the colon or other organs.
+Focus on the inner circular region captured by the endoscopic camera. Identify only small, raised, and rounded growths (polyps) with smooth and well-defined edges. These polyps typically appear brighter or redder than the surrounding tissue and have a soft, protruding texture. Ignore any flat areas, reflections, or irregular shapes that do not match these criteria. Avoid including empty or uniformly colored regions without distinct structures.
 Multiple polyps may exist in one image.
 """
 
+
 def highlight_mask_on_image(image, mask_np, color=(0, 200, 0), alpha=0.7):
     """Tô màu cho khu vực được xác định bởi mask lên hình ảnh gốc."""
     image_np = np.array(image)
@@ -25,24 +25,70 @@ def highlight_mask_on_image(image, mask_np, color=(0, 200, 0), alpha=0.7):
     highlighted_image[mask_indices] = (1 - alpha) * image_np[mask_indices] + alpha * np.array(color)
 
     return Image.fromarray(highlighted_image.astype(np.uint8))
+def display_masks_np(masks_np):
+    """
+    Hiển thị từng mask từ danh sách masks_np.
+
+    Parameters:
+        masks_np (list): Danh sách các mask dưới dạng numpy array.
+
+    Returns:
+        None: Hiển thị từng mask.
+    """
+    import matplotlib.pyplot as plt
+
+    print(f"Number of masks to display: {len(masks_np)}")
+    for i, mask_np in enumerate(masks_np):
+        print(f"Mask {i + 1} statistics: min={np.min(mask_np)}, max={np.max(mask_np)}, sum={np.sum(mask_np)}")
+
+        if np.sum(mask_np) > 0:  # Chỉ hiển thị mask nếu không rỗng
+            plt.imshow(mask_np, cmap="gray")
+            plt.title(f"Mask {i + 1}")
+            plt.axis("off")
+            plt.show()
+        else:
+            print(f"Mask {i + 1} is empty (all zeros).")
 
 def main(image):
+    # Convert image to RGB
+    image_pil = image.convert("RGB")
+
+    # Get prediction results from the model
+    res = model.predict([image_pil], [text_prompt])
+
+    # Check if the result has the expected structure
+    if not isinstance(res, list) or len(res) == 0 or not isinstance(res[0], dict):
+        raise ValueError("Unexpected result structure from model.predict")
+
+    # Extract masks, boxes, phrases, and logits from the result
+    masks = res[0].get("masks", [])
+    boxes = res[0].get("boxes", [])
+    phrases = res[0].get("phrases", [])
+    logits = res[0].get("scores", [])
+
+    # Handle the case where no masks are detected
+    if len(masks) == 0:
+        print(f"No objects of the '{text_prompt}' prompt detected in the image.")
+        return image_pil  # Return the original image if no masks are detected
 
-      image_pil = image.convert("RGB")
-      masks, boxes, phrases, logits = model.predict(image_pil, text_prompt)
+    # Convert masks to numpy arrays (if not already numpy arrays)
+    masks_np = [mask.squeeze() if isinstance(mask, np.ndarray) else mask.cpu().numpy() for mask in masks]
 
-      if len(masks) == 1:
-          print(f"No objects of the '{text_prompt}' prompt detected in the image.")
-          return image_pil  # Trả về ảnh gốc nếu không phát hiện ra mask
-      masks_np = [mask.squeeze().cpu().numpy() for mask in masks]
+    # Skip the first mask and process only masks from the second one onwards
+    if len(masks_np) > 1:
+        masks_np = masks_np[1:]  # Bỏ mask đầu tiên
 
-      if len(masks_np) > 1:
-          combined_mask = np.sum(masks_np[1:], axis=0) > 0  
-      else:
-          combined_mask = masks_np[0]  
-      highlighted_image = highlight_mask_on_image(image_pil, combined_mask)
+    # Combine masks if there are multiple masks
+    if len(masks_np) > 1:
+        combined_mask = np.sum(masks_np, axis=0) > 0
+    else:
+        combined_mask = masks_np[0]
 
-      return highlighted_image  
+    # Highlight the combined mask on the original image
+    highlighted_image = highlight_mask_on_image(image_pil, combined_mask)
+    #display_masks_np(masks_np)  # Display the remaining masks
+    return highlighted_image  # Return the highlighted image
+    
 def create_polyb():
     demo = gr.Interface(
         fn=main,