diff --git a/README.md b/README.md index ce5a381e..6ca1f254 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,10 @@ Some examples are listed below. You can find more in the directory of each model ![messi](./models/human_segmentation_pphumanseg/example_outputs/messi.jpg) +### Image Segmentation with [EfficientSAM](./models/image_segmentation_efficientsam/) + +![sam_present](./models/image_segmentation_efficientsam/example_outputs/sam_present.gif) + ### License Plate Detection with [LPD_YuNet](./models/license_plate_detection_yunet/) ![license plate detection](./models/license_plate_detection_yunet/example_outputs/lpd_yunet_demo.gif) diff --git a/models/__init__.py b/models/__init__.py index 1af41b7f..158e7687 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -20,6 +20,7 @@ from .facial_expression_recognition.facial_fer_model import FacialExpressionRecog from .object_tracking_vittrack.vittrack import VitTrack from .text_detection_ppocr.ppocr_det import PPOCRDet +from .image_segmentation_efficientsam.efficientSAM import EfficientSAM class ModuleRegistery: def __init__(self, name): @@ -94,3 +95,4 @@ def register(self, item): MODELS.register(FacialExpressionRecog) MODELS.register(VitTrack) MODELS.register(PPOCRDet) +MODELS.register(EfficientSAM) \ No newline at end of file diff --git a/models/image_segmentation_efficientsam/LICENSE b/models/image_segmentation_efficientsam/LICENSE new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/models/image_segmentation_efficientsam/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/models/image_segmentation_efficientsam/README.md b/models/image_segmentation_efficientsam/README.md new file mode 100644 index 00000000..d62a0a5d --- /dev/null +++ b/models/image_segmentation_efficientsam/README.md @@ -0,0 +1,52 @@ +# image_segmentation_efficientsam + +EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything + +Notes: +- The current implementation of the EfficientSAM demo uses the EfficientSAM-Ti model, which is specifically tailored for scenarios requiring higher speed and lightweight. +- image_segmentation_efficientsam_ti_2024may.onnx(supports only single point infering) + - MD5 value: 117d6a6cac60039a20b399cc133c2a60 + - SHA-256 value: e3957d2cd1422855f350aa7b044f47f5b3eafada64b5904ed330b696229e2943 +- image_segmentation_efficientsam_ti_2025april.onnx + - MD5 value: f23cecbb344547c960c933ff454536a3 + - SHA-256 value: 4eb496e0a7259d435b49b66faf1754aa45a5c382a34558ddda9a8c6fe5915d77 +- image_segmentation_efficientsam_ti_2025april_int8.onnx + - MD5 value: a1164f44b0495b82e9807c7256e95a50 + - SHA-256 value: 5ecc8d59a2802c32246e68553e1cf8ce74cf74ba707b84f206eb9181ff774b4e + + +## Demo + +### Python +Run the following command to try the demo: + +```shell +python demo.py --input /path/to/image +``` + +**Click** to select foreground points, **drag** to use box to select and **long press** to select background points on the object you wish to segment in the displayed image. After clicking the **Enter**, the segmentation result will be shown in a new window. Clicking the **Backspace** to clear all the prompts. + +## Result + +Here are some of the sample results that were observed using the model: + +![test1_res.jpg](./example_outputs/example1.png) +![test2_res.jpg](./example_outputs/example2.png) + +Video inference result: + +![sam_present.gif](./example_outputs/sam_present.gif) + +## Model metrics: + +## License + +All files in this directory are licensed under [Apache 2.0 License](./LICENSE). + +#### Contributor Details + +## Reference + +- https://arxiv.org/abs/2312.00863 +- https://github.com/yformer/EfficientSAM +- https://github.com/facebookresearch/segment-anything \ No newline at end of file diff --git a/models/image_segmentation_efficientsam/demo.py b/models/image_segmentation_efficientsam/demo.py new file mode 100644 index 00000000..a1410caf --- /dev/null +++ b/models/image_segmentation_efficientsam/demo.py @@ -0,0 +1,240 @@ +import argparse +import numpy as np +import cv2 as cv +from efficientSAM import EfficientSAM + +# Check OpenCV version +assert cv.__version__ >= "4.10.0", \ + "Please install latest opencv-python to try this demo: python3 -m pip install --upgrade opencv-python" + +# Valid combinations of backends and targets +backend_target_pairs = [ + [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_TARGET_CPU], + [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA], + [cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16], + [cv.dnn.DNN_BACKEND_TIMVX, cv.dnn.DNN_TARGET_NPU], + [cv.dnn.DNN_BACKEND_CANN, cv.dnn.DNN_TARGET_NPU] +] + +parser = argparse.ArgumentParser(description='EfficientSAM Demo') +parser.add_argument('--input', '-i', type=str, + help='Set input path to a certain image.') +parser.add_argument('--model', '-m', type=str, default='image_segmentation_efficientsam_ti_2024may.onnx', + help='Set model path, defaults to image_segmentation_efficientsam_ti_2024may.onnx.') +parser.add_argument('--backend_target', '-bt', type=int, default=0, + help='''Choose one of the backend-target pair to run this demo: + {:d}: (default) OpenCV implementation + CPU, + {:d}: CUDA + GPU (CUDA), + {:d}: CUDA + GPU (CUDA FP16), + {:d}: TIM-VX + NPU, + {:d}: CANN + NPU + '''.format(*[x for x in range(len(backend_target_pairs))])) +parser.add_argument('--save', '-s', action='store_true', + help='Specify to save a file with results. Invalid in case of camera input.') +args = parser.parse_args() + +# Global configuration +WINDOW_SIZE = (800, 600) # Fixed window size (width, height) +MAX_POINTS = 6 # Maximum allowed points +points = [] # Store clicked coordinates (original image scale) +labels = [] # Point labels (-1: useless, 0: background, 1: foreground, 2: top-left, 3: bottom right) +backend_point = [] +rectangle = False +current_img = None + +def visualize(image, result): + """ + Visualize the inference result on the input image. + + Args: + image (np.ndarray): The input image. + result (np.ndarray): The inference result. + + Returns: + vis_result (np.ndarray): The visualized result. + """ + # get image and mask + vis_result = np.copy(image) + mask = np.copy(result) + # change mask to binary image + t, binary = cv.threshold(mask, 127, 255, cv.THRESH_BINARY) + assert set(np.unique(binary)) <= {0, 255}, "The mask must be a binary image" + # enhance red channel to make the segmentation more obviously + enhancement_factor = 1.8 + red_channel = vis_result[:, :, 2] + # update the channel + red_channel = np.where(binary == 255, np.minimum(red_channel * enhancement_factor, 255), red_channel) + vis_result[:, :, 2] = red_channel + + # draw borders + contours, hierarchy = cv.findContours(binary, cv.RETR_LIST, cv.CHAIN_APPROX_TC89_L1) + cv.drawContours(vis_result, contours, contourIdx = -1, color = (255,255,255), thickness=2) + return vis_result + +def select(event, x, y, flags, param): + """Handle mouse events with coordinate conversion""" + global points, labels, backend_point, rectangle, current_img + orig_img = param['original_img'] + image_window = param['image_window'] + + if event == cv.EVENT_LBUTTONDOWN: + param['mouse_down_time'] = cv.getTickCount() + backend_point = [x, y] + + elif event == cv.EVENT_MOUSEMOVE: + if rectangle == True: + rectangle_change_img = current_img.copy() + cv.rectangle(rectangle_change_img, (backend_point[0], backend_point[1]), (x, y), (255,0,0) , 2) + cv.imshow(image_window, rectangle_change_img) + elif len(backend_point) != 0: + rectangle = True + + + elif event == cv.EVENT_LBUTTONUP: + if len(points) >= MAX_POINTS: + print(f"Maximum points reached ({MAX_POINTS})") + return + + if rectangle == False: + duration = (cv.getTickCount() - param['mouse_down_time'])/cv.getTickFrequency() + label = -1 if duration > 0.5 else 1 # Long press = background + + points.append([backend_point[0], backend_point[1]]) + labels.append(label) + print(f"Added {['background','foreground','background'][label]} point {backend_point}") + else: + if len(points) + 1 >= MAX_POINTS: + print(f"Points reached ({MAX_POINTS}), could not add box") + return + point_leftup = [] + point_rightdown = [] + if x > backend_point[0] or y > backend_point[1]: + point_leftup.extend(backend_point) + point_rightdown.extend([x,y]) + else: + point_leftup.extend([x,y]) + point_rightdown.extend(backend_point) + points.append(point_leftup) + points.append(point_rightdown) + print(f"Added box from {point_leftup} to {point_rightdown}") + labels.append(2) + labels.append(3) + rectangle = False + backend_point.clear() + + marked_img = orig_img.copy() + top_left = None + for (px, py), lbl in zip(points, labels): + if lbl == -1: + cv.circle(marked_img, (px, py), 5, (0, 0, 255), -1) + elif lbl == 1: + cv.circle(marked_img, (px, py), 5, (0, 255, 0), -1) + elif lbl == 2: + top_left = (px, py) + elif lbl == 3: + bottom_right = (px, py) + cv.rectangle(marked_img, top_left, bottom_right, (255,0,0) , 2) + cv.imshow(image_window, marked_img) + current_img = marked_img.copy() + + +if __name__ == '__main__': + backend_id = backend_target_pairs[args.backend_target][0] + target_id = backend_target_pairs[args.backend_target][1] + # Load the EfficientSAM model + model = EfficientSAM(modelPath=args.model) + + if args.input is not None: + # Read image + image = cv.imread(args.input) + if image is None: + print('Could not open or find the image:', args.input) + exit(0) + # create window + image_window = "Origin image" + cv.namedWindow(image_window, cv.WINDOW_NORMAL) + # change window size + rate = 1 + rate1 = 1 + rate2 = 1 + if(image.shape[1]>WINDOW_SIZE[0]): + rate1 = WINDOW_SIZE[0]/image.shape[1] + if(image.shape[0]>WINDOW_SIZE[1]): + rate2 = WINDOW_SIZE[1]/image.shape[0] + rate = min(rate1, rate2) + # width, height + WINDOW_SIZE = (int(image.shape[1] * rate), int(image.shape[0] * rate)) + cv.resizeWindow(image_window, WINDOW_SIZE[0], WINDOW_SIZE[1]) + # put the window on the left of the screen + cv.moveWindow(image_window, 50, 100) + # set listener to record user's click point + param = { + 'original_img': image, + 'mouse_down_time': 0, + 'image_window' : image_window + } + cv.setMouseCallback(image_window, select, param) + # tips in the terminal + print("Click — Select foreground point\n" + "Long press — Select background point\n" + "Drag — Create selection box\n" + "Enter — Infer\n" + "Backspace — Clear the prompts") + # show image + cv.imshow(image_window, image) + current_img = image.copy() + # create window to show visualized result + vis_image = image.copy() + segmentation_window = "Segment result" + cv.namedWindow(segmentation_window, cv.WINDOW_NORMAL) + cv.resizeWindow(segmentation_window, WINDOW_SIZE[0], WINDOW_SIZE[1]) + cv.moveWindow(segmentation_window, WINDOW_SIZE[0]+51, 100) + cv.imshow(segmentation_window, vis_image) + # waiting for click + while True: + # Check window status + # if click × to close the image window then ending + if (cv.getWindowProperty(image_window, cv.WND_PROP_VISIBLE) < 1 or + cv.getWindowProperty(segmentation_window, cv.WND_PROP_VISIBLE) < 1): + break + + # Handle keyboard input + key = cv.waitKey(1) + + # receive enter + if key == 13: + + vis_image = image.copy() + cv.putText(vis_image, "infering...", + (50, vis_image.shape[0]//2), + cv.FONT_HERSHEY_SIMPLEX, 10, (255,255,255), 5) + cv.imshow(segmentation_window, vis_image) + + result = model.infer(image=image, points=points, labels=labels) + if len(result) == 0: + print("clear and select points again!") + else: + vis_result = visualize(image, result) + + cv.imshow(segmentation_window, vis_result) + elif key == 8: # ASCII for Backspace + points.clear() + labels.clear() + backend_point = [] + rectangle = False + current_img = image + print("poins clear up") + cv.imshow(image_window, image) + + cv.destroyAllWindows() + + # Save results if save is true + if args.save: + cv.imwrite('./example_outputs/vis_result.jpg', vis_result) + cv.imwrite("./example_outputs/mask.jpg", result) + print('vis_result.jpg and mask.jpg are saved to ./example_outputs/') + + else: + print('Set input path to a certain image.') + pass + diff --git a/models/image_segmentation_efficientsam/efficientSAM.py b/models/image_segmentation_efficientsam/efficientSAM.py new file mode 100644 index 00000000..81077029 --- /dev/null +++ b/models/image_segmentation_efficientsam/efficientSAM.py @@ -0,0 +1,136 @@ +import numpy as np +import cv2 as cv + +class EfficientSAM: + def __init__(self, modelPath, backendId=0, targetId=0): + self._modelPath = modelPath + self._backendId = backendId + self._targetId = targetId + + self._model = cv.dnn.readNet(self._modelPath) + self._model.setPreferableBackend(self._backendId) + self._model.setPreferableTarget(self._targetId) + # 3 inputs + self._inputNames = ["batched_images", "batched_point_coords", "batched_point_labels"] + + self._outputNames = ['output_masks', 'iou_predictions'] # actual output layer name + self._currentInputSize = None + self._inputSize = [1024, 1024] # input size for the model + self._maxPointNums = 6 + self._frontGroundPoints = [] + self._backGroundPoints = [] + self._labels = [] + + @property + def name(self): + return self.__class__.__name__ + + def setBackendAndTarget(self, backendId, targetId): + self._backendId = backendId + self._targetId = targetId + self._model.setPreferableBackend(self._backendId) + self._model.setPreferableTarget(self._targetId) + + def _preprocess(self, image, points, labels): + + image = cv.cvtColor(image, cv.COLOR_BGR2RGB) + # record the input image size, (width, height) + self._currentInputSize = (image.shape[1], image.shape[0]) + + image = cv.resize(image, self._inputSize) + + image = image.astype(np.float32, copy=False) / 255.0 + + image_blob = cv.dnn.blobFromImage(image) + + points = np.array(points, dtype=np.float32) + labels = np.array(labels, dtype=np.float32) + assert points.shape[0] <= self._maxPointNums, f"Max input points number: {self._maxPointNums}" + assert points.shape[0] == labels.shape[0] + + frontGroundPoints = [] + backGroundPoints = [] + inputLabels = [] + for i in range(len(points)): + if labels[i] == -1: + backGroundPoints.append(points[i]) + else: + frontGroundPoints.append(points[i]) + inputLabels.append(labels[i]) + self._backGroundPoints = np.uint32(backGroundPoints) + print("input:") + print(" back: ", self._backGroundPoints) + print(" front: ", frontGroundPoints) + print(" label: ", inputLabels) + + # convert points to (1024*1024) size space + for p in frontGroundPoints: + p[0] = np.float32(p[0] * self._inputSize[0]/self._currentInputSize[0]) + p[1] = np.float32(p[1] * self._inputSize[1]/self._currentInputSize[1]) + + if len(frontGroundPoints) > self._maxPointNums: + return "no" + + pad_num = self._maxPointNums - len(frontGroundPoints) + self._frontGroundPoints = np.vstack([frontGroundPoints, np.zeros((pad_num, 2), dtype=np.float32)]) + inputLabels_arr = np.array(inputLabels, dtype=np.float32).reshape(-1, 1) + self._labels = np.vstack([inputLabels_arr, np.full((pad_num, 1), -1, dtype=np.float32)]) + + points_blob = np.array([[self._frontGroundPoints]]) + + labels_blob = np.array([[self._labels]]) + + return image_blob, points_blob, labels_blob + + def infer(self, image, points, labels): + # Preprocess + imageBlob, pointsBlob, labelsBlob = self._preprocess(image, points, labels) + # Forward + self._model.setInput(imageBlob, self._inputNames[0]) + self._model.setInput(pointsBlob, self._inputNames[1]) + self._model.setInput(labelsBlob, self._inputNames[2]) + print("infering...") + outputs = self._model.forward(self._outputNames) + outputBlob, outputIou = outputs[0], outputs[1] + # Postprocess + results = self._postprocess(outputBlob, outputIou) + print("done") + return results + + def _postprocess(self, outputBlob, outputIou): + # The masks are already sorted by their predicted IOUs. + # The first dimension is the batch size (we have a single image. so it is 1). + # The second dimension is the number of masks we want to generate + # The third dimension is the number of candidate masks output by the model. + masks = outputBlob[0, 0, :, :, :] >= 0 + ious = outputIou[0, 0, :] + + # sorted by ious + sorted_indices = np.argsort(ious)[::-1] + sorted_masks = masks[sorted_indices] + + # sorted by area + # mask_areas = np.sum(masks, axis=(1, 2)) + # sorted_indices = np.argsort(mask_areas) + # sorted_masks = masks[sorted_indices] + + masks_uint8 = (sorted_masks * 255).astype(np.uint8) + + # change to real image size + resized_masks = [ + cv.resize(mask, dsize=self._currentInputSize, + interpolation=cv.INTER_NEAREST) + for mask in masks_uint8 + ] + + # background mask don't need + for mask in resized_masks: + contains_bg = any( + mask[y, x] if (0 <= x < mask.shape[1] and 0 <= y < mask.shape[0]) + else False + for (x, y) in self._backGroundPoints + ) + if not contains_bg: + return mask + + return resized_masks[0] diff --git a/models/image_segmentation_efficientsam/example_outputs/example1.png b/models/image_segmentation_efficientsam/example_outputs/example1.png new file mode 100644 index 00000000..c20d7834 --- /dev/null +++ b/models/image_segmentation_efficientsam/example_outputs/example1.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:70065831fb12915dc5a3b4641019bc152a89d6d5be1887bdf7ada432a04e63c5 +size 1993654 diff --git a/models/image_segmentation_efficientsam/example_outputs/example2.png b/models/image_segmentation_efficientsam/example_outputs/example2.png new file mode 100644 index 00000000..3b0cb955 --- /dev/null +++ b/models/image_segmentation_efficientsam/example_outputs/example2.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dfe6860d701b8b707a96d69b6bfc33fd05167168fbb46594f6377ad4e9c1733e +size 1917383 diff --git a/models/image_segmentation_efficientsam/example_outputs/sam_present.gif b/models/image_segmentation_efficientsam/example_outputs/sam_present.gif new file mode 100644 index 00000000..403a2817 --- /dev/null +++ b/models/image_segmentation_efficientsam/example_outputs/sam_present.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ab75c654d4368d1f4762fc71af35c02b6f0a3e21dca4530d22f92fff4134890c +size 103918 diff --git a/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april.onnx b/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april.onnx new file mode 100644 index 00000000..2bf444b7 --- /dev/null +++ b/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4eb496e0a7259d435b49b66faf1754aa45a5c382a34558ddda9a8c6fe5915d77 +size 48312857 diff --git a/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april_int8.onnx b/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april_int8.onnx new file mode 100644 index 00000000..8f7b6907 --- /dev/null +++ b/models/image_segmentation_efficientsam/image_segmentation_efficientsam_ti_2025april_int8.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ecc8d59a2802c32246e68553e1cf8ce74cf74ba707b84f206eb9181ff774b4e +size 20479928