TF port of the Segment Anything Model (SAM) (#22970)
* First commit * Add auto-translation with GPT-4 * make fixup * Add a functional layernorm for TF * Add all the auxiliary imports etc. * Add the extra processor and tests * rebase to main * Add all the needed fixes to the GPT code * make fixup * Make convolutions channels-last so they run on CPU * make fixup * Fix final issues * Fix other models affected by test change * Clarify comment on the sparse_prompt_embeddings check * Refactor functional_layernorm, use shape_list in place of .shape in some places * Remove deprecated torch-alike code * Update tests/models/sam/test_modeling_tf_sam.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * Update tests/models/sam/test_modeling_tf_sam.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * Refactor processor with common methods and separated private methods * make fixup * Quietly delete the file that didn't do anything (sorry Sylvain) * Refactor the processor tests into one file * make fixup * Clean up some unnecessary indirection * Fix TF mask postprocessing * Add more processor equivalence tests * Refactor generate_crop_boxes to use framework-neutral np code * Make the serving output correctly conditional * Fix error message line length * Use dict keys rather than indices internally in both TF and PT SAM call/forward * Return dicts internally in the call/forward methods * Revert changes to common tests and just override check_pt_tf_outputs * Revert changes to other model tests * Clarify comments for functional layernorm * Add missing transpose from PT code * Removed unused copied from in PT code * Remove overrides for tests that don't exist in TF * Fix transpose and update tests for PT and TF to check pred_masks * Add training flag * Update tests to use TF checkpoints * Update index.mdx * Add missing cross-test decorator * Remove optional extra asterisks * Revert return_dict changes in PT code * Update src/transformers/models/sam/modeling_tf_sam.py Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> * Remove None return annotations on init methods * Update tests/models/sam/test_processor_sam.py Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> * Fix input_boxes shapes * make fixup --------- Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com> Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
This commit is contained in:
Matt
@@ -17,8 +17,14 @@ import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from transformers.testing_utils import require_torch, require_torchvision, require_vision
|
||||
from transformers.utils import is_torch_available, is_vision_available
|
||||
from transformers.testing_utils import (
|
||||
is_pt_tf_cross_test,
|
||||
require_tf,
|
||||
require_torch,
|
||||
require_torchvision,
|
||||
require_vision,
|
||||
)
|
||||
from transformers.utils import is_tf_available, is_torch_available, is_vision_available
|
||||
|
||||
|
||||
if is_vision_available():
|
||||
@@ -29,6 +35,9 @@ if is_vision_available():
|
||||
if is_torch_available():
|
||||
import torch
|
||||
|
||||
if is_tf_available():
|
||||
import tensorflow as tf
|
||||
|
||||
|
||||
@require_vision
|
||||
@require_torchvision
|
||||
@@ -110,3 +119,158 @@ class SamProcessorTest(unittest.TestCase):
|
||||
dummy_masks = [[1, 0], [0, 1]]
|
||||
with self.assertRaises(ValueError):
|
||||
masks = processor.post_process_masks(dummy_masks, np.array(original_sizes), np.array(reshaped_input_size))
|
||||
|
||||
|
||||
@require_vision
|
||||
@require_tf
|
||||
class TFSamProcessorTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
image_processor = SamImageProcessor()
|
||||
processor = SamProcessor(image_processor)
|
||||
processor.save_pretrained(self.tmpdirname)
|
||||
|
||||
def get_image_processor(self, **kwargs):
|
||||
return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).image_processor
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmpdirname)
|
||||
|
||||
def prepare_image_inputs(self):
|
||||
"""This function prepares a list of PIL images, or a list of numpy arrays if one specifies numpify=True,
|
||||
or a list of PyTorch tensors if one specifies torchify=True.
|
||||
"""
|
||||
|
||||
image_inputs = [np.random.randint(255, size=(3, 30, 400), dtype=np.uint8)]
|
||||
|
||||
image_inputs = [Image.fromarray(np.moveaxis(x, 0, -1)) for x in image_inputs]
|
||||
|
||||
return image_inputs
|
||||
|
||||
def test_save_load_pretrained_additional_features(self):
|
||||
processor = SamProcessor(image_processor=self.get_image_processor())
|
||||
processor.save_pretrained(self.tmpdirname)
|
||||
|
||||
image_processor_add_kwargs = self.get_image_processor(do_normalize=False, padding_value=1.0)
|
||||
|
||||
processor = SamProcessor.from_pretrained(self.tmpdirname, do_normalize=False, padding_value=1.0)
|
||||
|
||||
self.assertEqual(processor.image_processor.to_json_string(), image_processor_add_kwargs.to_json_string())
|
||||
self.assertIsInstance(processor.image_processor, SamImageProcessor)
|
||||
|
||||
def test_image_processor(self):
|
||||
image_processor = self.get_image_processor()
|
||||
|
||||
processor = SamProcessor(image_processor=image_processor)
|
||||
|
||||
image_input = self.prepare_image_inputs()
|
||||
|
||||
input_feat_extract = image_processor(image_input, return_tensors="np")
|
||||
input_processor = processor(images=image_input, return_tensors="np")
|
||||
|
||||
input_feat_extract.pop("original_sizes") # pop original_sizes as it is popped in the processor
|
||||
input_feat_extract.pop("reshaped_input_sizes") # pop reshaped_input_sizes as it is popped in the processor
|
||||
|
||||
for key in input_feat_extract.keys():
|
||||
self.assertAlmostEqual(input_feat_extract[key].sum(), input_processor[key].sum(), delta=1e-2)
|
||||
|
||||
@require_tf
|
||||
def test_post_process_masks(self):
|
||||
image_processor = self.get_image_processor()
|
||||
|
||||
processor = SamProcessor(image_processor=image_processor)
|
||||
dummy_masks = [tf.ones((1, 3, 5, 5))]
|
||||
|
||||
original_sizes = [[1764, 2646]]
|
||||
|
||||
reshaped_input_size = [[683, 1024]]
|
||||
masks = processor.post_process_masks(dummy_masks, original_sizes, reshaped_input_size, return_tensors="tf")
|
||||
self.assertEqual(masks[0].shape, (1, 3, 1764, 2646))
|
||||
|
||||
masks = processor.post_process_masks(
|
||||
dummy_masks,
|
||||
tf.convert_to_tensor(original_sizes),
|
||||
tf.convert_to_tensor(reshaped_input_size),
|
||||
return_tensors="tf",
|
||||
)
|
||||
self.assertEqual(masks[0].shape, (1, 3, 1764, 2646))
|
||||
|
||||
# should also work with np
|
||||
dummy_masks = [np.ones((1, 3, 5, 5))]
|
||||
masks = processor.post_process_masks(
|
||||
dummy_masks, np.array(original_sizes), np.array(reshaped_input_size), return_tensors="tf"
|
||||
)
|
||||
|
||||
self.assertEqual(masks[0].shape, (1, 3, 1764, 2646))
|
||||
|
||||
dummy_masks = [[1, 0], [0, 1]]
|
||||
with self.assertRaises(tf.errors.InvalidArgumentError):
|
||||
masks = processor.post_process_masks(
|
||||
dummy_masks, np.array(original_sizes), np.array(reshaped_input_size), return_tensors="tf"
|
||||
)
|
||||
|
||||
|
||||
@require_vision
|
||||
@require_torchvision
|
||||
class SamProcessorEquivalenceTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdirname = tempfile.mkdtemp()
|
||||
image_processor = SamImageProcessor()
|
||||
processor = SamProcessor(image_processor)
|
||||
processor.save_pretrained(self.tmpdirname)
|
||||
|
||||
def get_image_processor(self, **kwargs):
|
||||
return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).image_processor
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.tmpdirname)
|
||||
|
||||
def prepare_image_inputs(self):
|
||||
"""This function prepares a list of PIL images, or a list of numpy arrays if one specifies numpify=True,
|
||||
or a list of PyTorch tensors if one specifies torchify=True.
|
||||
"""
|
||||
|
||||
image_inputs = [np.random.randint(255, size=(3, 30, 400), dtype=np.uint8)]
|
||||
|
||||
image_inputs = [Image.fromarray(np.moveaxis(x, 0, -1)) for x in image_inputs]
|
||||
|
||||
return image_inputs
|
||||
|
||||
@is_pt_tf_cross_test
|
||||
def test_post_process_masks_equivalence(self):
|
||||
image_processor = self.get_image_processor()
|
||||
|
||||
processor = SamProcessor(image_processor=image_processor)
|
||||
dummy_masks = np.random.randint(0, 2, size=(1, 3, 5, 5)).astype(np.float32)
|
||||
tf_dummy_masks = [tf.convert_to_tensor(dummy_masks)]
|
||||
pt_dummy_masks = [torch.tensor(dummy_masks)]
|
||||
|
||||
original_sizes = [[1764, 2646]]
|
||||
|
||||
reshaped_input_size = [[683, 1024]]
|
||||
tf_masks = processor.post_process_masks(
|
||||
tf_dummy_masks, original_sizes, reshaped_input_size, return_tensors="tf"
|
||||
)
|
||||
pt_masks = processor.post_process_masks(
|
||||
pt_dummy_masks, original_sizes, reshaped_input_size, return_tensors="pt"
|
||||
)
|
||||
|
||||
self.assertTrue(np.all(tf_masks[0].numpy() == pt_masks[0].numpy()))
|
||||
|
||||
@is_pt_tf_cross_test
|
||||
def test_image_processor_equivalence(self):
|
||||
image_processor = self.get_image_processor()
|
||||
|
||||
processor = SamProcessor(image_processor=image_processor)
|
||||
|
||||
image_input = self.prepare_image_inputs()
|
||||
|
||||
pt_input_feat_extract = image_processor(image_input, return_tensors="pt")["pixel_values"].numpy()
|
||||
pt_input_processor = processor(images=image_input, return_tensors="pt")["pixel_values"].numpy()
|
||||
|
||||
tf_input_feat_extract = image_processor(image_input, return_tensors="tf")["pixel_values"].numpy()
|
||||
tf_input_processor = processor(images=image_input, return_tensors="tf")["pixel_values"].numpy()
|
||||
|
||||
self.assertTrue(np.allclose(pt_input_feat_extract, pt_input_processor))
|
||||
self.assertTrue(np.allclose(pt_input_feat_extract, tf_input_feat_extract))
|
||||
self.assertTrue(np.allclose(pt_input_feat_extract, tf_input_processor))
|
||||
|
||||
Reference in New Issue
Block a user