Cryo-SAM
Author: James Rodgers, Aditya Balu
Organization: Iowa State University - Department of Electrical and Computer Engineering
Email: jmsrdgrs@iastate.edu

Uses Sam for the detection of squares, holes, and particles in cryo-em.

Usage:

from cryo_sam import Csam

Model checkpoint location, several types downloadable from https://github.com/facebookresearch/segment-anything#model-checkpoints
testing done with the below model and checkpoint information:

    model_name = "vit_h"
    checkpoint = "sam_model/sam_vit_h_4b8939.pth" 
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Example construction:
    detector = Csam(model_name, checkpoint, device)

Csam.hole_finder(numpy_image_array)
    Parameter: numpy image array
    Returns: list of square coordinates (x, y, w, h)
Csam.square_finder(numpy_image_array):
    Parameter: numpy image array
    Returns: list of square coordinates (x, y, w, h)

Both of the above returns will need to be changed in the future since extra data 
is to be returned be the exact nature of that isn't determined yet. It will probably return a class 
containing the data in the near future.

Csam.visualize_detections(image, image_scale_percent, box_prompts = None, point_prompts = None, 
                          show_boxes = True, show_point_prompts = None, show_box_prompts, show_masks = True):
    Parameters:
            image: numpy image array, numpy array of points, numpy array of bounding box, bool, bool, bool
            image_scale_percent: int representing scaled percent i.e 60 = 60% of original size
            point_prompts: numpy array of points. [[x, y], [x, y], [x, y], ...]
            box_prompts: numpy array of rectangle boxes [[x1, y1, x2, y2], [x1, y1, x2, y2], ...]
            show_boxes: bool, show bounding boxes generated by the model
            show_masks: bool, show masks generated by the model
            show_point_prompts: bool, show point prompts sent to the model
            show_box_prompts: bool, show bounding box prompts sent to the model
    Returns: None

    Example usage:

    test_image = np.load("/mnt/big_data_2/pncc/annotated_data/160002/box4_g11/data_031022/mmm_0.npy")
    test_image = cv2.cvtColor(test_image, cv2.COLOR_BGR2RGB)

    print(detector.hole_finder(test_image))
    detector.visualize_detections(test_image, image_scale_percent, show_boxes=True)

    image_scale_percent = 40

    This function will display the image with the detections overlaid on top of it. It is currently not specialized to hole or square.
    This will be changed in the future, this is mostly for testing purposes.
    The image_scale_percent is the percent of the original image size that the image will be displayed at. 
    This is useful for large images that would be too large to display on a screen and to improve render speed. 
    The default is 50% of the original size. Will probably add a save image option in the future.


