change foundation model from grounding dino to clip

src/app.py

@@ -31,7 +31,7 @@ lab_df_input = gr.Dataframe(
     headers=["Object", "New A", "New B"],
     datatype=["str", "number", "number"],
     col_count=(3, "fixed"),
-    label="Target Objects and New Settings",
+    label="Target Objects and New Settings.(0-255 -- 128 = Neutral)",
     type="array",
 )
 
@@ -78,15 +78,15 @@ change_color_objects_lab_tool = gr.Interface(
     examples=[
         [
             "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_1.jpg",
-            [["pants", 128, 1]],
+            [["pants", 112, 128]],
         ],
         [
             "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_4.jpg",
-            [["desk", 15, 0.5], ["left cup", 40, 1.1]],
+            [["desk", 166, 193]],
         ],
         [
             "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_5.jpg",
-            [["suits", 60, 1.5], ["pants", 10, 0.8]],
+            [["suits coat", 110, 133]],
         ],
     ],
 )
@@ -117,6 +117,16 @@ privacy_preserve_tool = gr.Interface(
             "license plate.",
             10,
         ],
+        [
+            "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_8.jpg",
+            "face.",
+            15,
+        ],
+        [
+            "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_6.jpg",
+            "face.",
+            20,
+        ],
     ],
 )
 
@@ -135,21 +145,24 @@ remove_background_tool = gr.Interface(
         [
             "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_6.jpg",
         ],
+        [
+            "https://raw.githubusercontent.com/mahan-ym/ImageAlfred/main/src/assets/examples/test_8.jpg",
+        ],
     ],
 )
 
 demo = gr.TabbedInterface(
     [
-        change_color_objects_hsv_tool,
-        change_color_objects_lab_tool,
         privacy_preserve_tool,
         remove_background_tool,
+        change_color_objects_hsv_tool,
+        change_color_objects_lab_tool,
     ],
     [
-        "Change Color Objects HSV",
-        "Change Color Objects LAB",
         "Privacy Preserving Tool",
         "Remove Background Tool",
+        "Change Color Objects HSV",
+        "Change Color Objects LAB",
     ],
     title=title,
     theme=gr.themes.Default(

src/modal_app.py

@@ -5,7 +5,6 @@ import cv2
 import modal
 import numpy as np
 from PIL import Image
-from rapidfuzz import process
 
 app = modal.App("ImageAlfred")
 
@@ -30,14 +29,16 @@ image = (
             "TORCH_HOME": TORCH_HOME,
         }
     )
-    .apt_install("git")
+    .apt_install(
+        "git",
+    )
     .pip_install(
         "huggingface-hub",
         "hf_transfer",
         "Pillow",
         "numpy",
+        "transformers",
         "opencv-contrib-python-headless",
-        "RapidFuzz",
         gpu="A10G",
     )
     .pip_install(
@@ -46,10 +47,8 @@ image = (
         index_url="https://download.pytorch.org/whl/cu124",
         gpu="A10G",
     )
-    .pip_install(
-        "git+https://github.com/luca-medeiros/lang-segment-anything.git",
-        gpu="A10G",
-    )
+    .pip_install("git+https://github.com/openai/CLIP.git", gpu="A10G")
+    .pip_install("git+https://github.com/facebookresearch/sam2.git", gpu="A10G")
     .pip_install(
         "git+https://github.com/PramaLLC/BEN2.git#egg=ben2",
         gpu="A10G",
@@ -58,43 +57,180 @@ image = (
 
 
 @app.function(
+    image=image,
+    gpu="A10G",
+    volumes={volume_path: volume},
+    timeout=60 * 3,
+)
+def prompt_segment(
+    image_pil: Image.Image,
+    prompts: list[str],
+) -> list[dict]:
+    clip_results = clip.remote(image_pil, prompts)
+
+    if not clip_results:
+        print("No boxes returned from CLIP.")
+        return None
+
+    boxes = np.array(clip_results["boxes"])
+
+    sam_result_masks, sam_result_scores = sam2.remote(image_pil=image_pil, boxes=boxes)
+
+    print(f"sam_result_mask {sam_result_masks}")
+
+    if not sam_result_masks.any():
+        print("No masks or scores returned from SAM2.")
+        return None
+
+    if sam_result_masks.ndim == 3:
+        # If the masks are in 3D, we need to convert them to 4D
+        sam_result_masks = [sam_result_masks]
+
+    results = {
+        "labels": clip_results["labels"],
+        "boxes": boxes,
+        "clip_scores": clip_results["scores"],
+        "sam_masking_scores": sam_result_scores,
+        "masks": sam_result_masks,
+    }
+    return results
+
+
+@app.function(
+    image=image,
     gpu="A10G",
+    volumes={volume_path: volume},
+    timeout=60 * 3,
+)
+def sam2(image_pil: Image.Image, boxes: list[np.ndarray]) -> list[dict]:
+    import torch
+    from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        predictor.set_image(image_pil)
+        masks, scores, _ = predictor.predict(
+            point_coords=None,
+            point_labels=None,
+            box=boxes,
+            multimask_output=False,
+        )
+    return masks, scores
+
+
+@app.function(
     image=image,
+    gpu="A10G",
     volumes={volume_path: volume},
-    # min_containers=1,
     timeout=60 * 3,
 )
-def lang_sam_segment(
+def clip(
     image_pil: Image.Image,
-    prompt: str,
-    box_threshold=0.3,
-    text_threshold=0.25,
-) -> list:
-    """Segments an image using LangSAM based on a text prompt.
-    This function uses LangSAM to segment objects in the image based on the provided prompt.
-    """  # noqa: E501
-    from lang_sam import LangSAM  # type: ignore
-
-    model = LangSAM(sam_type="sam2.1_hiera_large")
-    langsam_results = model.predict(
-        images_pil=[image_pil],
-        texts_prompt=[prompt],
-        box_threshold=box_threshold,
-        text_threshold=text_threshold,
+    prompts: list[str],
+) -> list[dict]:
+    """
+    returns:
+        dict with keys each are lists:
+            - labels: str, the prompt used for the prediction
+            - scores: float, confidence score of the prediction
+            - boxes: np.array representing bounding box coordinates
+    """
+
+    from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
+    import torch
+
+    processor = CLIPSegProcessor.from_pretrained(
+        "CIDAS/clipseg-rd64-refined",
+        use_fast=True,
     )
-    if len(langsam_results[0]["labels"]) == 0:
-        print("No masks found for the given prompt.")
-        return None
+    model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
-    print(f"found {len(langsam_results[0]['labels'])} masks for prompt: {prompt}")
-    print("labels:", langsam_results[0]["labels"])
-    print("scores:", langsam_results[0]["scores"])
-    print(
-        "masks scores:",
-        langsam_results[0].get("mask_scores", "No mask scores available"),
-    )  # noqa: E501
+    # Get original image dimensions
+    orig_width, orig_height = image_pil.size
 
-    return langsam_results
+    inputs = processor(
+        text=prompts,
+        images=[image_pil] * len(prompts),
+        padding="max_length",
+        return_tensors="pt",
+    )
+    # predict
+    with torch.no_grad():
+        outputs = model(**inputs)
+    preds = outputs.logits.unsqueeze(1)
+
+    # Get the dimensions of the prediction output
+    pred_height, pred_width = preds.shape[-2:]
+
+    # Calculate scaling factors
+    width_scale = orig_width / pred_width
+    height_scale = orig_height / pred_height
+
+    labels = []
+    scores = []
+    boxes = []
+
+    # Process each prediction to find bounding boxes in high probability regions
+    for i, prompt in enumerate(prompts):
+        # Apply sigmoid to get probability map
+        pred_tensor = torch.sigmoid(preds[i][0])
+        # Convert tensor to numpy array
+        pred_np = pred_tensor.cpu().numpy()
+
+        # Convert to uint8 for OpenCV processing
+        heatmap = (pred_np * 255).astype(np.uint8)
+
+        # Apply threshold to find high probability regions
+        _, binary = cv2.threshold(heatmap, 127, 255, cv2.THRESH_BINARY)
+
+        # Find contours in thresholded image
+        contours, _ = cv2.findContours(
+            binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        # Process each contour to get bounding boxes
+        for contour in contours:
+            # Skip very small contours that might be noise
+            if cv2.contourArea(contour) < 100:  # Minimum area threshold
+                continue
+
+            # Get bounding box coordinates in prediction space
+            x, y, w, h = cv2.boundingRect(contour)
+
+            # Scale coordinates to original image dimensions
+            x_orig = int(x * width_scale)
+            y_orig = int(y * height_scale)
+            w_orig = int(w * width_scale)
+            h_orig = int(h * height_scale)
+
+            # Calculate confidence score based on average probability in the region
+            mask = np.zeros_like(pred_np)
+            cv2.drawContours(mask, [contour], 0, 1, -1)
+            confidence = float(np.mean(pred_np[mask == 1]))
+
+            labels.append(prompt)
+            scores.append(confidence)
+            boxes.append(
+                np.array(
+                    [
+                        x_orig,
+                        y_orig,
+                        x_orig + w_orig,
+                        y_orig + h_orig,
+                    ]
+                )
+            )
+
+    if labels == []:
+        return None
+
+    results = {
+        "labels": labels,
+        "scores": scores,
+        "boxes": boxes,
+    }
+    return results
 
 
 @app.function(
@@ -128,14 +264,16 @@ def change_image_objects_hsv(
             "targets_config must be a list of lists, each containing [target_name, hue, saturation_scale]."  # noqa: E501
         )
     print("Change image objects hsv targets config:", targets_config)
-    prompts = ". ".join(target[0] for target in targets_config)
+    prompts = [target[0].strip() for target in targets_config]
 
-    langsam_results = lang_sam_segment.remote(image_pil=image_pil, prompt=prompts)
-    if not langsam_results:
+    prompt_segment_results = prompt_segment.remote(
+        image_pil=image_pil,
+        prompts=prompts,
+    )
+    if not prompt_segment_results:
         return image_pil
-    input_labels = [target[0] for target in targets_config]
-    output_labels = langsam_results[0]["labels"]
-    scores = langsam_results[0]["scores"]
+
+    output_labels = prompt_segment_results["labels"]
 
     img_array = np.array(image_pil)
     img_hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV).astype(np.float32)
@@ -144,13 +282,14 @@ def change_image_objects_hsv(
         if not label or label == "":
             print("Skipping empty label.")
             continue
-        input_label, score, _ = process.extractOne(label, input_labels)
-        input_label_idx = input_labels.index(input_label)
-
+        if label not in prompts:
+            print(f"Label '{label}' not found in prompts. Skipping.")
+            continue
+        input_label_idx = prompts.index(label)
         target_rgb = targets_config[input_label_idx][1:]
         target_hsv = cv2.cvtColor(np.uint8([[target_rgb]]), cv2.COLOR_RGB2HSV)[0][0]
 
-        mask = langsam_results[0]["masks"][idx].astype(bool)
+        mask = prompt_segment_results["masks"][idx][0].astype(bool)
         h, s, v = cv2.split(img_hsv)
         # Convert all channels to float32 for consistent processing
         h = h.astype(np.float32)
@@ -168,9 +307,9 @@ def change_image_objects_hsv(
         scale_s = target_s / mean_s if mean_s > 0 else 1.0
         scale_v = target_v / mean_v if mean_v > 0 else 1.0
 
-        scale_s = np.clip(scale_s, 0.8, 1.2)  
+        scale_s = np.clip(scale_s, 0.8, 1.2)
         scale_v = np.clip(scale_v, 0.8, 1.2)
-        
+
         # Apply changes only in mask
         h[mask] = target_hue
         s = s.astype(np.float32)
@@ -224,18 +363,16 @@ def change_image_objects_lab(
 
     print("change image objects lab targets config:", targets_config)
 
-    prompts = ". ".join(target[0] for target in targets_config)
+    prompts = [target[0].strip() for target in targets_config]
 
-    langsam_results = lang_sam_segment.remote(
+    prompt_segment_results = prompt_segment.remote(
         image_pil=image_pil,
-        prompt=prompts,
+        prompts=prompts,
     )
-    if not langsam_results:
+    if not prompt_segment_results:
         return image_pil
 
-    input_labels = [target[0] for target in targets_config]
-    output_labels = langsam_results[0]["labels"]
-    scores = langsam_results[0]["scores"]
+    output_labels = prompt_segment_results["labels"]
 
     img_array = np.array(image_pil)
     img_lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2Lab).astype(np.float32)
@@ -244,13 +381,17 @@ def change_image_objects_lab(
         if not label or label == "":
             print("Skipping empty label.")
             continue
-        input_label, score, _ = process.extractOne(label, input_labels)
-        input_label_idx = input_labels.index(input_label)
+
+        if label not in prompts:
+            print(f"Label '{label}' not found in prompts. Skipping.")
+            continue
+
+        input_label_idx = prompts.index(label)
 
         new_a = targets_config[input_label_idx][1]
         new_b = targets_config[input_label_idx][2]
 
-        mask = langsam_results[0]["masks"][idx]
+        mask = prompt_segment_results["masks"][idx][0]
         mask_bool = mask.astype(bool)
 
         img_lab[mask_bool, 1] = new_a
@@ -298,49 +439,46 @@ def apply_mosaic_with_bool_mask(
 )
 def preserve_privacy(
     image_pil: Image.Image,
-    prompt: str,
+    prompts: str,
     privacy_strength: int = 15,
 ) -> Image.Image:
     """
     Preserves privacy in an image by applying a mosaic effect to specified objects.
     """
-    print(f"Preserving privacy for prompt: {prompt} with strength {privacy_strength}")
-
-    langsam_results = lang_sam_segment.remote(
+    print(f"Preserving privacy for prompt: {prompts} with strength {privacy_strength}")
+    if isinstance(prompts, str):
+        prompts = [prompt.strip() for prompt in prompts.split(".")]
+        print(f"Parsed prompts: {prompts}")
+    prompt_segment_results = prompt_segment.remote(
         image_pil=image_pil,
-        prompt=prompt,
-        box_threshold=0.35,
-        text_threshold=0.40,
+        prompts=prompts,
     )
-    if not langsam_results:
+    if not prompt_segment_results:
         return image_pil
 
     img_array = np.array(image_pil)
 
-    for result in langsam_results:
-        print(f"result: {result}")
-
-        for i, mask in enumerate(result["masks"]):
-            if "mask_scores" in result:
-                if (
-                    hasattr(result["mask_scores"], "shape")
-                    and result["mask_scores"].ndim > 0
-                ):
-                    mask_score = result["mask_scores"][i]
-                else:
-                    mask_score = result["mask_scores"]
-            if mask_score < 0.6:
-                print(f"Skipping mask {i + 1}/{len(result['masks'])} -> low score.")
-                continue
-            print(
-                f"Processing mask {i + 1}/{len(result['masks'])} Mask score: {mask_score}"  # noqa: E501
-            )
+    for i, mask in enumerate(prompt_segment_results["masks"]):
+        mask_bool = mask[0].astype(bool)
 
-            mask_bool = mask.astype(bool)
+        # Create kernel for morphological operations
+        kernel_size = 100
+        kernel = np.ones((kernel_size, kernel_size), np.uint8)
 
-            img_array = apply_mosaic_with_bool_mask.remote(
-                img_array, mask_bool, privacy_strength
-            )
+        # Convert bool mask to uint8 for OpenCV operations
+        mask_uint8 = mask_bool.astype(np.uint8) * 255
+
+        # Apply dilation to slightly expand the mask area
+        mask_uint8 = cv2.dilate(mask_uint8, kernel, iterations=2)
+        # Optional: Apply erosion again to refine the mask
+        mask_uint8 = cv2.erode(mask_uint8, kernel, iterations=2)
+
+        # Convert back to boolean mask
+        mask_bool = mask_uint8 > 127
+
+        img_array = apply_mosaic_with_bool_mask.remote(
+            img_array, mask_bool, privacy_strength
+        )
 
     output_image_pil = Image.fromarray(img_array)
 
@@ -354,14 +492,14 @@ def preserve_privacy(
     timeout=60 * 2,
 )
 def remove_background(image_pil: Image.Image) -> Image.Image:
-    from ben2 import BEN_Base
-    import torch
+    import torch  # type: ignore
+    from ben2 import BEN_Base  # type: ignore
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
     print("type of image_pil:", type(image_pil))
     model = BEN_Base.from_pretrained("PramaLLC/BEN2")
-    model.to(device).eval()
+    model.to(device).eval()  # todo check if this should be outside the function
 
     output_image = model.inference(
         image_pil,

src/tools.py

@@ -23,7 +23,7 @@ def remove_background(
     if not input_img:
         raise gr.Error("Input image cannot be None or empty.")
 
-    func = modal.Function.from_name("ImageAlfred", "remove_background")
+    func = modal.Function.from_name(modal_app_name, "remove_background")
     output_pil = func.remote(
         image_pil=input_img,
     )
@@ -67,10 +67,10 @@ def privacy_preserve_image(
     if not valid_pattern.match(input_prompt):
         raise gr.Error("Input prompt must contain only letters, spaces, and dots.")
 
-    func = modal.Function.from_name("ImageAlfred", "preserve_privacy")
+    func = modal.Function.from_name(modal_app_name, "preserve_privacy")
     output_pil = func.remote(
         image_pil=input_img,
-        prompt=input_prompt,
+        prompts=input_prompt,
         privacy_strength=privacy_strength,
     )
 
@@ -136,14 +136,14 @@ def change_color_objects_hsv(
             raise gr.Error("Red must be an integer.")
         if item[1] < 0 or item[1] > 255:
             raise gr.Error("Red must be in the range [0, 255]")
-        
+
         try:
             item[2] = int(item[2])
         except ValueError:
             raise gr.Error("Green must be an integer.")
         if item[2] < 0 or item[2] > 255:
             raise gr.Error("Green must be in the range [0, 255]")
-        
+
         try:
             item[3] = int(item[3])
         except ValueError:
@@ -153,7 +153,7 @@ def change_color_objects_hsv(
 
     print("after processing input:", user_input)
 
-    func = modal.Function.from_name("ImageAlfred", "change_image_objects_hsv")
+    func = modal.Function.from_name(modal_app_name, "change_image_objects_hsv")
     output_pil = func.remote(image_pil=input_img, targets_config=user_input)
 
     if output_pil is None:
@@ -248,7 +248,7 @@ def change_color_objects_lab(
                 raise gr.Error("new B must be in the range [0, 255]")
 
     print("after processing input:", user_input)
-    func = modal.Function.from_name("ImageAlfred", "change_image_objects_lab")
+    func = modal.Function.from_name(modal_app_name, "change_image_objects_lab")
     output_pil = func.remote(image_pil=input_img, targets_config=user_input)
     if output_pil is None:
         raise ValueError("Received None from modal remote function.")