import warnings
import numpy as np
import requests
from PIL import Image
import gradio as gr
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.
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)

    if mask_np.shape != image_np.shape[:2]:
        raise ValueError("Kích thước của mask không khớp với kích thước của hình ảnh")

    highlighted_image = image_np.copy()

    mask_indices = mask_np > 0  # Chỉ lấy các vùng có mask
    highlighted_image[mask_indices] = (1 - alpha) * image_np[mask_indices] + alpha * np.array(color)

    return Image.fromarray(highlighted_image.astype(np.uint8))

def main(image):

      image_pil = image.convert("RGB")
      masks, boxes, phrases, logits = model.predict(image_pil, text_prompt)

      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]

      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)

      return highlighted_image  
def create_polyb():
    demo = gr.Interface(
        fn=main,
        inputs=gr.Image(type="pil"),
        outputs=gr.Image(type="pil"),
        title="Highlight Polyps",
        description="Tải lên một hình ảnh và phát hiện polyp với LangSAM."
    )  
    return demo