Spaces:

DanielFD
/

IPD-app

Runtime error

IPD-app / f_segment_img.py

Upload 4 files

ebc51c9 over 3 years ago

6.66 kB

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

	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_photoarea_frame_max > area_object > area_photoarea_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