Revolutionizing Medical Imaging with SAM and W&B

!nvidia-smi 

import os
HOME = os.getcwd()
print("HOME:", HOME)

!pip install -q 'git+https://github.com/facebookresearch/segment-anything.git'
!pip install -q jupyter_bbox_widget roboflow dataclasses-json supervision

import os
!mkdir -p {HOME}/weights
!wget -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P {HOME}/weights
CHECKPOINT_PATH = os.path.join(HOME, "weights", "sam_vit_h_4b8939.pth")
print(CHECKPOINT_PATH, "; exist:", os.path.isfile(CHECKPOINT_PATH))

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor


sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)

%cd {HOME}


import roboflow
from roboflow import Roboflow


roboflow.login()


rf = Roboflow()


project = rf.workspace("hashira-fhxpj").project("mri-brain-tumor")
dataset = project.version(1).download("coco")

import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Union, Optional
from dataclasses_json import dataclass_json
from supervision import Detections
@dataclass_json
@dataclass
class COCOCategory:
   id: int
   name: str
   supercategory: str

@dataclass_json
@dataclass
class COCOImage:
   id: int
   width: int
   height: int
   file_name: str
   license: int
   date_captured: str
   coco_url: Optional[str] = None
   flickr_url: Optional[str] = None

@dataclass_json
@dataclass
class COCOAnnotation:
   id: int
   image_id: int
   category_id: int
   segmentation: List[List[float]]
   area: float
   bbox: Tuple[float, float, float, float]
   iscrowd: int

@dataclass_json
@dataclass
class COCOLicense:
   id: int
   name: str
   url: str

@dataclass_json
@dataclass
class COCOJson:
   images: List[COCOImage]
   annotations: List[COCOAnnotation]
   categories: List[COCOCategory]
   licenses: List[COCOLicense]




def load_coco_json(json_file: str) -> COCOJson:
   import json


   with open(json_file, "r") as f:
       json_data = json.load(f)


   return COCOJson.from_dict(json_data)

class COCOJsonUtility:
   @staticmethod
   def get_annotations_by_image_id(coco_data: COCOJson, image_id: int) -> List[COCOAnnotation]:
       return [annotation for annotation in coco_data.annotations if annotation.image_id == image_id]

@staticmethod
   def get_annotations_by_image_path(coco_data: COCOJson, image_path: str) -> Optional[List[COCOAnnotation]]:
       image = COCOJsonUtility.get_image_by_path(coco_data, image_path)
       if image:
           return COCOJsonUtility.get_annotations_by_image_id(coco_data, image.id)
       else:
           return None


@staticmethod
   def get_image_by_path(coco_data: COCOJson, image_path: str) -> Optional[COCOImage]:
       for image in coco_data.images:
           if image.file_name == image_path:
               return image
       return None

@staticmethod
   def annotations2detections(annotations: List[COCOAnnotation]) -> Detections:
       class_id, xyxy = [], []


       for annotation in annotations:
           x_min, y_min, width, height = annotation.bbox
           class_id.append(annotation.category_id)
           xyxy.append([
               x_min,
               y_min,
               x_min + width,
               y_min + height
           ])


       return Detections(
           xyxy=np.array(xyxy, dtype=int),
           class_id=np.array(class_id, dtype=int)
       )

import os


DATA_SET_SUBDIRECTORY = "test"
ANNOTATIONS_FILE_NAME = "_annotations.coco.json"
IMAGES_DIRECTORY_PATH = os.path.join(dataset.location, DATA_SET_SUBDIRECTORY)
ANNOTATIONS_FILE_PATH = os.path.join(dataset.location, DATA_SET_SUBDIRECTORY, ANNOTATIONS_FILE_NAME)
coco_data = load_coco_json(json_file=ANNOTATIONS_FILE_PATH)


CLASSES = [
   category.name
   for category
   in coco_data.categories
   if category.supercategory != 'none'
]


IMAGES = [
   image.file_name
   for image
   in coco_data.images
]

import random
random.seed(10)
mask_predictor = SamPredictor(sam)

for image_name in IMAGES:
   image_path = os.path.join(dataset.location, DATA_SET_SUBDIRECTORY, image_name)


   # Load dataset annotations
   annotations = COCOJsonUtility.get_annotations_by_image_path(coco_data=coco_data, image_path=image_name)
   ground_truth = COCOJsonUtility.annotations2detections(annotations=annotations)


   # Adjust class ids
   ground_truth.class_id = ground_truth.class_id - 1


   # Load image
   image_bgr = cv2.imread(image_path)
   image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

# Initiate annotator
   bounding_box_annotator = sv.BoundingBoxAnnotator(color=sv.Color.red(), color_lookup=sv.ColorLookup.INDEX)
   mask_annotator = sv.MaskAnnotator(color=sv.Color.red(), color_lookup=sv.ColorLookup.INDEX)


   # Annotate ground truth
   annotated_frame_ground_truth = bounding_box_annotator.annotate(scene=image_bgr.copy(), detections=ground_truth)

# Run SAM inference
   mask_predictor.set_image(image_rgb)
   masks, scores, logits = mask_predictor.predict(
       box=ground_truth.xyxy[0],
       multimask_output=True
   )


   detections = sv.Detections(
       xyxy=sv.mask_to_xyxy(masks=masks),
       mask=masks
   )
   detections = detections[detections.area == np.max(detections.area)]


   annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)


   sv.plot_images_grid(
       images=[annotated_frame_ground_truth, annotated_image],
       grid_size=(1, 2),
       titles=['source image', 'segmented image']
   )

!pip install wandb
import wandb
from PIL import Image


# Initialize W&B
wandb.login()
run = wandb.init(project="segment-anything", name="sam-multiple-images")


def mask_to_image(mask):
   """Convert a binary mask to an RGB image."""
   mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
   mask = mask.astype(np.uint8) * 255
   return Image.fromarray(mask)

# Loop through multiple images
for image_name in IMAGES:
   image_path = os.path.join(dataset.location, DATA_SET_SUBDIRECTORY, image_name)


   # Load dataset annotations
   annotations = COCOJsonUtility.get_annotations_by_image_path(coco_data=coco_data, image_path=image_name)
   ground_truth = COCOJsonUtility.annotations2detections(annotations=annotations)


   # Adjust class ids
   ground_truth.class_id = ground_truth.class_id - 1

# Load image
   image_bgr = cv2.imread(image_path)
   image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)


   # Run SAM inference
   mask_predictor.set_image(image_rgb)
   masks, scores, logits = mask_predictor.predict(
       box=ground_truth.xyxy[0],
       multimask_output=True
   )


   detections = sv.Detections(
       xyxy=sv.mask_to_xyxy(masks=masks),
       mask=masks
   )
   detections = detections[detections.area == np.max(detections.area)]

# Annotate images
   bounding_box_annotator = sv.BoundingBoxAnnotator(color=sv.Color.red(), color_lookup=sv.ColorLookup.INDEX)
   mask_annotator = sv.MaskAnnotator(color=sv.Color.red(), color_lookup=sv.ColorLookup.INDEX)
   annotated_frame_ground_truth = bounding_box_annotator.annotate(scene=image_bgr.copy(), detections=ground_truth)
   annotated_image = mask_annotator.annotate(scene=image_bgr.copy(), detections=detections)


   # Convert annotated images to PIL Images for logging
   source_image_pil = Image.fromarray(cv2.cvtColor(annotated_frame_ground_truth, cv2.COLOR_BGR2RGB))
   segmented_image_pil = Image.fromarray(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))

# Log images and masks to W&B
   wandb.log({
       "source_image": wandb.Image(source_image_pil, caption=f"Source Image with Ground Truth ({image_name})"),
       "segmented_image": wandb.Image(segmented_image_pil, caption=f"Segmented Image by SAM ({image_name})"),
       "masks": [wandb.Image(mask_to_image(mask), caption=f"Mask {i} for {image_name}") for i, mask in enumerate(masks)],
       "scores": {f"score_{i}": score for i, score in enumerate(scores)}  # Log scores for each mask
   })


# Finish the W&B run
wandb.finish()