Add `automatic-mask-generation` pipeline for Segment Anything Model (SAM) #22840

ArthurZucker · 2023-04-18T17:05:41Z

What does this PR do?

This need the SAM model + rebasing once merged

from transformers import pipeline
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import time

generator = pipeline("automatic-mask-generation", device = 0)
image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"

dog_url = "/home/arthur_huggingface_co/transformers/Arthur/dog.jpg"
raw_image = Image.open(dog_url).convert("RGB")

start = time.time()
outputs = generator(raw_image, points_per_batch = 256, pred_iou_thresh=1)
print(f"point_batch_size : {256}, {time.time() - start}")

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    

plt.imshow(np.array(raw_image))
ax = plt.gca()
for mask in outputs["masks"]:
    show_mask(mask, ax=ax, random_color=True)
plt.axis("off")
plt.show()

plt.savefig("dog_results_2.png")

HuggingFaceDocBuilderDev · 2023-04-18T17:27:24Z

The documentation is not available anymore as the PR was closed or merged.

sgugger

Thanks for your PR! Left some initial comments. Also why the "automatic" in "automatic-mask-generation"? "mask-generation" is clear enough no?

src/transformers/pipelines/__init__.py

src/transformers/pipelines/automatic_mask_generation.py

Narsil

Overall looks good.

I think in general SAM could be just image-segmentation, but there seems to be a lot of specificities here with a lot of custom code, so making it standalone is ok for me now.

Custom code is marked as private so we can move later. And we could always make this pipeline be an alias of image-segmentation

src/transformers/pipelines/automatic_mask_generation.py

tests/pipelines/test_pipelines_automatic_mask_generation.py

Co-authored-by: Nicolas Patry <patry.nicolas@gmail.com>

Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com>

amyeroberts

Nice - super exciting to see this pipeline!

+1 to all of Sylvain's comments about docstrings and variable names.

For the post processing that happens, I think most of the functionality should sit with the image processor. And the image processor should have dedicated methods for this. Other model's processors similarly filter boxes and perform RLE conversion.

For quite a few of the methods, I found it quite confusing what the functions were doing and the objects being handled. A lot of this should be resolved with Sylvain's suggestions. For some, splitting into more atomic functions and having consistent types e.g. not having to handle many different mask shapes would also help.

amyeroberts · 2023-04-19T13:29:05Z

src/transformers/pipelines/automatic_mask_generation.py

@@ -0,0 +1,615 @@
+import math


General question (for @Narsil ?) - how come we don't have copyright headers for pipeline files?

I don't know.

We can add them.
When I create new files I tend to copy/paste from something else, I might have missed some.

If copyright headers are important, shouldn't we have some sort of lint for them ?

src/transformers/pipelines/automatic_mask_generation.py

amyeroberts · 2023-04-19T15:46:04Z

src/transformers/pipelines/mask_generation.py

+    masks = masks > mask_threshold
+    converted_boxes = _batched_mask_to_box(masks)
+
+    keep_mask = ~_is_box_near_crop_edge(converted_boxes, cropped_box_image, [0, 0, original_width, original_height])


What is cropped_box_image (type and what does it represent)?

src/transformers/pipelines/mask_generation.py

…to add-mg-pipeline

…nsformers into HEAD

…nsformers into add-mg-pipeline

sgugger

Would love for @amyeroberts to have a second look, but LGTM! Thanks!

amyeroberts

Nice update - structure looking a lit tidier!

A few main points:

There's still a lot of issues with the docstrings: missing, incomplete, wrong which need to be updated
A white_pixels check should still be part of the tests
Overall pipeline code looks good 👍 Just a few general nits there regarding argument values
I'm a bit concerned about the processing code. There's a lot of assumptions about the image types and shapes which I'm not sure are always correct. For a first pass of the postprocessing, we don't have to make it compatible in all cases, but it should be double checked and assumptions about inputs stated in the docstrings or comments.
Is the mask_threshold value right?

amyeroberts · 2023-04-20T14:18:58Z

tests/pipelines/test_pipelines_mask_generation.py

+        self.assertEqual(
+            nested_simplify(new_outupt, decimals=4),
+            [
+                {'mask': {'hash': '115ad19f5f', 'shape': (480, 640)}, 'scores': 1.0444},


There should be a check here on the pixel value counts similar to the white_pixels in image segmentation. As mentioned before, if a single pixel changes the hash is completely different and without additional information test debugging is a lot harder. Checking the model output, the masks are binary, as so white_pixels check should be easy to add. If all values are 0, as before then this indicates an issue with the outputs.

Agreed, I will let @ArthurZucker do this in a follow up PR befrore the next release as discussed offline!

amyeroberts · 2023-04-20T14:21:46Z

src/transformers/pipelines/mask_generation.py

+        all_boxes = []
+        for model_output in model_outputs:
+            all_scores.append(model_output.pop("iou_scores"))
+            all_masks.extend(model_output.pop("masks"))


extend difference here hasn't been addressed. See: #22840 (comment)

The reason behind that is that model_output.pop("masks") returns a list of masks, and post_process_for_mask_generation expects a single list of masks instead of nested lists. Therefore you need to call extend

src/transformers/pipelines/mask_generation.py

src/transformers/models/sam/image_processing_sam.py

amyeroberts · 2023-04-20T15:12:31Z

src/transformers/models/sam/image_processing_sam.py

+
+    def post_process_for_mask_generation(self, all_masks, all_scores, all_boxes, crops_nms_thresh):
+        """
+        Post processes mask that are automatically generated.


Docstring missing:

Information about the output - what do the post processed outputs look like and represent?

Information about the input arguments and their types

src/transformers/models/sam/image_processing_sam.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

amyeroberts

LGTM - thanks for iterating!

src/transformers/models/sam/image_processing_sam.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

younesbelkada · 2023-04-20T16:16:25Z

Thank you all for your reviews!

src/transformers/feature_extraction_utils.py

HuggingFaceDocBuilderDev · 2023-04-20T17:33:20Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

…SAM) (huggingface#22840) * cleanup * updates * more refactoring * make style * update inits * support other inputs in base * update based on review Co-authored-by: Nicolas Patry <patry.nicolas@gmail.com> * Update tests/pipelines/test_pipelines_automatic_mask_generation.py Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com> * update * fixup * TODO x and y to refactor, _h _w refactored here * update docstring * more nits * style on these * more doc fix * rename variables * update * updates * style * update * fix `_mask_to_rle_pytorch` * styling * fix ask to rle, wrong outputs * add device arg * update * more updates, fix tets * udpate * update docstrings * styling * fixup * add notebook on the docs * update orginal sizes * fix docstring * updat condition on point_per-batch * updates tests * fix CI test * extend is required, append does not work! * fixup * fix CI tests * whit pixels left * address doc comments * fix doc * slow pipeline tests * update auto init * add revision * make fixup * update p!ipoeline tag when calling tests * alphabeitcal order in inits * fix copies * last style nits * Apply suggestions from code review Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * reformat docstring * more reformat * address most of the comments * Update src/transformers/pipelines/mask_generation.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * final refactor * Update src/transformers/models/sam/image_processing_sam.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * fixup and fix slow tests * revert --------- Co-authored-by: Nicolas Patry <patry.nicolas@gmail.com> Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com> Co-authored-by: younesbelkada <younesbelkada@gmail.com> Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

ArthurZucker force-pushed the add-mg-pipeline branch from 9a3b075 to 4eeb27d Compare April 18, 2023 18:03

cleanup

eb4dc26

ArthurZucker force-pushed the add-mg-pipeline branch from 4eeb27d to eb4dc26 Compare April 18, 2023 18:05

younesbelkada mentioned this pull request Apr 19, 2023

Add Segment Anything Model (SAM) #22654

Merged

ArthurZucker added 4 commits April 19, 2023 09:43

updates

f3f345b

more refactoring

94fd59d

make style

6bbb106

update inits

52675a6

ArthurZucker requested review from Narsil and sgugger April 19, 2023 10:07

ArthurZucker marked this pull request as ready for review April 19, 2023 11:51

sgugger reviewed Apr 19, 2023

View reviewed changes

Narsil reviewed Apr 19, 2023

View reviewed changes

src/transformers/pipelines/automatic_mask_generation.py Outdated Show resolved Hide resolved

tests/pipelines/test_pipelines_automatic_mask_generation.py Outdated Show resolved Hide resolved

sgugger requested a review from amyeroberts April 19, 2023 13:15

ArthurZucker and others added 6 commits April 19, 2023 13:43

support other inputs in base

3901d13

update based on review

df77716

Co-authored-by: Nicolas Patry <patry.nicolas@gmail.com>

Update tests/pipelines/test_pipelines_automatic_mask_generation.py

9d9e70b

Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com>

update

0a7f702

fixup

d57ed60

TODO x and y to refactor, _h _w refactored here

2d11d6f

amyeroberts reviewed Apr 19, 2023

View reviewed changes

ArthurZucker and others added 5 commits April 19, 2023 16:42

update docstring

8950318

more nits

c68453a

style on these

e8b45b3

more doc fix

ef08e55

rename variables

d9027ed

younesbelkada changed the title ~~[WIP] Add mg pipeline~~ [WIP] Add automatic-mask-generation pipeline for Segment Anything Model (SAM) Apr 19, 2023

ArthurZucker added 2 commits April 19, 2023 19:02

update

0e507b8

Merge branch 'main' of https://github.com/huggingface/transformers in…

3b6f918

…to add-mg-pipeline

younesbelkada added 3 commits April 20, 2023 13:30

add revision

b9e4272

Merge branch 'add-mg-pipeline' of https://github.com/ArthurZucker/tra…

749b94f

…nsformers into HEAD

make fixup

d5eb3f9

younesbelkada requested a review from sgugger April 20, 2023 13:31

ArthurZucker added 2 commits April 20, 2023 13:31

update p!ipoeline tag when calling tests

e103ef2

alphabeitcal order in inits

ed4e727

younesbelkada changed the title ~~[WIP] Add automatic-mask-generation pipeline for Segment Anything Model (SAM)~~ Add automatic-mask-generation pipeline for Segment Anything Model (SAM) Apr 20, 2023

ArthurZucker added 3 commits April 20, 2023 13:32

fix copies

3c4f4b8

last style nits

eb9bd0c

Merge branch 'add-mg-pipeline' of https://github.com/ArthurZucker/tra…

d98f0e3

…nsformers into add-mg-pipeline

sgugger approved these changes Apr 20, 2023

View reviewed changes

amyeroberts reviewed Apr 20, 2023

View reviewed changes

younesbelkada and others added 6 commits April 20, 2023 17:34

Apply suggestions from code review

7a1f99a

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

reformat docstring

94a7af3

more reformat

7d8b860

address most of the comments

a59133a

Update src/transformers/pipelines/mask_generation.py

e970f30

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

final refactor

b5d56da

younesbelkada requested a review from amyeroberts April 20, 2023 16:09

amyeroberts approved these changes Apr 20, 2023

View reviewed changes

src/transformers/models/sam/image_processing_sam.py Outdated Show resolved Hide resolved

Update src/transformers/models/sam/image_processing_sam.py

32d6e80

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

younesbelkada added 2 commits April 20, 2023 16:39

fixup and fix slow tests

4149f36

Merge remote-tracking branch 'upstream/main' into HEAD

9a0c19d

younesbelkada reviewed Apr 20, 2023

View reviewed changes

src/transformers/feature_extraction_utils.py Outdated Show resolved Hide resolved

revert

3b4a594

younesbelkada merged commit f143037 into huggingface:main Apr 20, 2023

ArthurZucker deleted the add-mg-pipeline branch April 21, 2023 09:12

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `automatic-mask-generation` pipeline for Segment Anything Model (SAM) #22840

Add `automatic-mask-generation` pipeline for Segment Anything Model (SAM) #22840

Add automatic-mask-generation pipeline for Segment Anything Model (SAM) #22840

Add automatic-mask-generation pipeline for Segment Anything Model (SAM) #22840

Conversation

What does this PR do?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Add `automatic-mask-generation` pipeline for Segment Anything Model (SAM) #22840

Add `automatic-mask-generation` pipeline for Segment Anything Model (SAM) #22840