File size: 6,657 Bytes
ebc51c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
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 |