f_segment_img.py

from f_measurents import *
import cv2
import numpy as np
import matplotlib.pyplot as plt
import mediapipe as mp
from mediapipe.tasks.python import vision
import sys
import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
#
sys.path.insert(
    1, '/Users/danielfiuzadosil/Documents/GitHub_Repo/Bryant_Medical/eCommerce/App_IPD [Master]/ipd_app/src/modules')
#


def remove_background_img(img):
    THRESHOLD = 0.12
    # initialize mediapipe
    mp_selfie_segmentation = mp.solutions.selfie_segmentation
    selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation(
        model_selection=1)
    # get the result
    results = selfie_segmentation.process(img)
    # extract segmented mask
    mask = np.stack((results.segmentation_mask,) * 3, axis=-1) > THRESHOLD
    mask_binary = mask.astype(int)*255
    img_masked = img.copy()
    img_masked[mask_binary == 0] = 255
    #
    return img_masked


def remove_background_img_v2(image_path, MODEL_PATH='/Users/danielfiuzadosil/Documents/GitHub_Repo/Bryant_Medical/eCommerce/App_IPD [Master]/ipd_app/data/external/mediapipe_models/deeplabv3.tflite'):
    #
    THRESHOLD = 0.12
    #
    BG_COLOR = (0, 0, 0)  # black
    MASK_COLOR = (255, 255, 255)  # white
    #
    BaseOptions = mp.tasks.BaseOptions
    OutputType = vision.ImageSegmenterOptions.OutputType
    # Create the options that will be used for ImageSegmenter
    base_options = BaseOptions(model_asset_path=MODEL_PATH)
    options = vision.ImageSegmenterOptions(
        base_options=base_options, output_type=OutputType.CATEGORY_MASK)
    # Create the MediaPipe image file that will be segmented
    image = mp.Image.create_from_file(image_path)
    with vision.ImageSegmenter.create_from_options(options) as segmenter:
        # Retrieve the masks for the segmented image
        category_masks = segmenter.segment(image)
    # Generate solid color images for showing the output segmentation mask.
    image_data = image.numpy_view()
    fg_image = np.zeros(image_data.shape, dtype=np.uint8)
    fg_image[:] = MASK_COLOR
    bg_image = np.zeros(image_data.shape, dtype=np.uint8)
    bg_image[:] = BG_COLOR
    #
    condition = np.stack(
        (category_masks[0].numpy_view(),) * 3, axis=-1) > THRESHOLD
    mask_binary = np.where(condition, fg_image, bg_image)
    #
    img_masked = image_data.copy()
    img_masked[mask_binary == 0] = 255
    #
    return img_masked


def segment_frame_from_img(image, landmarks, sam):
    # Read image
    image_0 = image.copy()
    # Generate facial landmarks
    mp_face_mesh = mp.solutions.face_mesh
    face_mesh = mp_face_mesh.FaceMesh(
        static_image_mode=True,
        max_num_faces=1,
        refine_landmarks=True,
        min_detection_confidence=0.5)
    # Calculate facial landmaks and other data
    result, mesh_points = face_mesh_points(image)
    df_mesh_points = mesh_points_to_df(mesh_points)
    width_face_px, height_face_px = get_face_dimensions_px(
        mesh_points, landmarks)
    # Calculate rectangle where the frame will likely be (based on a reference landmarks)
    squareBoxEyes = mesh_points[landmarks['squareBoxEyes']]
    #
    width_rectangle = squareBoxEyes[2][0] - squareBoxEyes[0][0]
    tolerance_x = 0.2*width_rectangle
    height_rectangle = squareBoxEyes[2][1] - squareBoxEyes[0][1]
    tolerance_y = height_rectangle*0.2
    #
    x_start, y_start = [squareBoxEyes[0][0] -
                        tolerance_x, squareBoxEyes[0][1] - tolerance_y]
    x_end, y_end = [squareBoxEyes[2][0] + tolerance_x, squareBoxEyes[2][1]]
    # Cropped image to region where the frames will be located
    img_cropped = image_0[int(y_start):int(y_end), int(x_start):int(x_end)]
    # Use Meta's Segment Anything Model (SAM) to segment the frame
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(img_cropped)
    # Select the right object by defining the expected range of area occupied by the frame
    height, width, _ = img_cropped.shape
    area_photo = height*width
    area_frame_min = 0.2
    area_frame_max = 0.6
    # Iterate through the different masks and store the ones that fulfill the criteria
    masks_selection = []
    objects_segmented = []
    for i in range(len(masks)):
        mask = masks[i]
        area_object = mask['area']
        if area_photo*area_frame_max > area_object > area_photo*area_frame_min:
            masks_selection.append(mask)
            #
            mask_binary_temp = mask['segmentation'].astype(int)*255
            object_segmented = img_cropped.copy()
            object_segmented[mask_binary_temp == 0] = 255
            objects_segmented.append(object_segmented)

    return image, img_cropped, masks_selection, objects_segmented


def plot_sam_check_segmentation_frame(image, img_cropped, objects_segmented):
    #
    ax1 = plt.subplot(311)
    ax1.imshow(image)
    ax1.axis("Off")
    ax2 = plt.subplot(312)
    ax2.imshow(img_cropped)
    ax2.axis("Off")
    ax3 = plt.subplot(313)
    ax3.imshow(objects_segmented[0])
    fig = plt.gcf()
    return fig
    

def plot_sam_check_segmentation_frame_and_save(image, img_cropped, objects_segmented, output_folder, filepath):
    #
    ax1 = plt.subplot(311)
    ax1.imshow(image)
    ax1.axis("Off")
    ax2 = plt.subplot(312)
    ax2.imshow(img_cropped)
    ax2.axis("Off")
    ax3 = plt.subplot(313)
    ax3.imshow(objects_segmented[0])
    #
    plt.savefig(output_folder + os.path.basename(filepath),
                transparent=True, bbox_inches='tight')
    # plt.show()

def plot_segmented_object_with_bb(objects_segmented, masks_selection, image):
    image = cv2.rectangle(objects_segmented[0], masks_selection[0]['bbox'], [255,0,0], 4)
    plt.imshow(image)
    ax = plt.gca()
    # 
    img_height = image.shape[0]
    img_width = image.shape[1]
    # 
    major_ticks_height = np.arange(0, img_height, 100)
    minor_ticks_height = np.arange(0, img_height, 10)
    major_ticks_width = np.arange(0, img_width, 100)
    minor_ticks_width = np.arange(0, img_width, 10)
    # 
    ax.set_yticks(major_ticks_height)
    ax.set_yticks(minor_ticks_height, minor=True)
    ax.set_xticks(major_ticks_width)
    ax.set_xticks(minor_ticks_width, minor=True)
    # 
    ax.grid(which='both')
    ax.grid(which='minor', alpha=0.2)
    ax.grid(which='major', alpha=0.5)
    # 
    # plt.show()

def get_frame_width(masks_selection):
    frame_width = masks_selection[0]['bbox'][2]
    return frame_width

def ipd_calibration(ipd_px, frame_width_px, frame_width_mm):
    calibration_factor = frame_width_px/frame_width_mm
    ipd_mm = ipd_px/calibration_factor
    return ipd_mm