331 - Fine-tune Segment Anything Model (SAM) using custom data

This tutorial walks you through the process of fine-tuning a Segment Anything Model (SAM) using custom data.
Code from this video is available here: github.com/bnsreenu/python_fo...
What is SAM?
SAM is an image segmentation model developed by Meta AI. It was trained over 11 billion segmentation masks from millions of images. It is designed to take human prompts, in the form of points, bounding boxes or even a text prompt describing what should be segmented.
What are the key features of SAM?
Zero-shot generalization: SAM can be used to segment objects that it has never seen before, without the need for additional training.
Flexible prompting: SAM can be prompted with a variety of input, including points, boxes, and text descriptions.
Real-time mask computation: SAM can generate masks for objects in real time. This makes SAM ideal for applications where it is necessary to segment objects quickly, such as autonomous driving and robotics.
Ambiguity awareness: SAM is aware of the ambiguity of objects in images. This means that SAM can generate masks for objects even when they are partially occluded or overlapping with other objects.
How does SAM work?
SAM works by first encoding the image into a high-dimensional vector representation. The prompt is encoded into a separate vector representation. The two vector representations are then combined and passed to a mask decoder, which outputs a mask for the object specified by the prompt.
The image encoder is a vision transformer (ViT-H) model, which is a large language model that has been pre-trained on a massive dataset of images. The prompt encoder is a simple text encoder that converts the input prompt into a vector representation. The mask decoder is a lightweight transformer model that predicts the object mask from the image and prompt embeddings.
SAM paper: arxiv.org/pdf/2304.02643.pdf​
Link to the dataset used in this demonstration: www.epfl.ch/labs/cvlab/data/d...
Courtesy: EPFL
This code has been heavily adapted from this notebook but modified to work with a truly custom dataset where we have a bunch of images and binary masks. github.com/NielsRogge/Transfo...

[@kevian182]

Excellent tutorial Sreeni!!! 👏👏Thank you so much!!!

[@mith888]

Классное видео ! Спасибо за подробное объяснение!

[@user-ki3ls5th4i]

Awesome. Thanks for this detailed explanation. It helped me a lot as a starter practitioner of SAM.

[@DigitalSreeni]

Glad it helped!

[@dmitryutkin9864]

Thank you very much for such a wonderful tutorial!!!

[@user-wh7ld4jj6e]

Really great video. Thank you so much.

[@robosergTV]

this is gold, thanks

[@surajprasad8741]

Thank you sir, got clear understanding

[@philipplagrange314]

Great video, thank you! It would be interesting to know how to relate SAM to other models for additional classification! Could you possibly make a video about it?

[@llz-gp1db]

Nice video. Thanks for sharing!!!

[@timanb2491]

great job! thanks!

[@NicolaRomano]

Great video as always. I think the function to find bboxes might be improved to take care of the fact that you might have multiple objects in a patch (I guess you could do a simple watershed and then find min and max for each instance). Also I'm wondering if you could improve results by adding some heuristics to how you choose your grid points, for instance concentrating points in darker areas in this case?

[@AnusuyaT-gz5zc]

Your videos are so good.. please post a video on deep image prior.. Thanks

[@user-yj7th8lb9h]

Is there a way that we can use SAM for an image sequence? I'm trying to segment grains and pore area for small sand.

[@gabrielgcarvalho]

Great video, and great instructor. However... This get_bounding_box is not very good for multiple objects. Furthermore, I could not make it work for more than one bounding box as a prompt. Do you have an idea how to generalize it?

[@maheethabharadwaj8016]

Thank you so much for this incredible and praactical video. Is there a way to segment multiple different objects within the same model or does it need to be two separate? For example if i wanted to segment both mitochondria and lysosomes (and train a model to recognizes BOTH those things but as different things). would i need a separate SAM for mito vs lysosomes? Is there a way to do it that would be combined?

[@jerinantony007]

Hi, good content. How can we train overlapping case? Train with one box and it's segment mask at a time? Or can we train with all boxes at a time utilising three output channels?

[@djondle]

Thanks!

[@DigitalSreeni]

Thank you very much.

[@mahmoudma3n935]

Could you make a video on how to use the SAM image encoder only as a feature extractor and then use any other decoder to get the prediction mask?

[@hamzawi2752]

I was going through the same problem of drop_last=True. This is simply because if the last batch in your dataset contains only 1 training sample, you will get this error since batch normalization can be applied to one training sample. For instance, if the batch size is 2, and your training dataset is 101, in this case, you have 51 batches, the last batch contains only one training sample, and this absolutely will throw an error. You can generate this error and comment right here.

[@hik381]

Great video. If we have multiple objects in an image that we want to fine tune, should we create one mask for each image with all objects masked and having like multiple bboxes , or a separate mask for each object in the same image?

[@ultimaterocker4]

Hi did you ever figure this one out?

[@juliannad9879]

This is great thanks a lot ! However, since you deleted the images with empty masks, this means that this can work only for images where there are mitochondria. Could this be extended so that the model returns an empty mask when there is no mito ? (or other things for other applications)

[@youmustbenewhereguy]

How to finetune a multiclass segmentation label? How to make the prompt based on the label too?

[@I_A_D_L]

how to measure the masks created from the SAM model? Thank you very much!.

[@phoenix1799]

How to make a tif file for images and masks if I have custom data to train or is there any work around to train the model on custom data?

[@user-lx9xw6fb5q]

Hello Sir! I want to fine-tune my satellite datasets to delineate crop field parcels. But I am confused how to prepare masks for them. I want each crop parcel has different number (like instance segmentation). But it seems this tutorial provide for binary segmentation. How to solve this issue? Can you give me some advice to prepare masks datasets?

[@barryjuait]

And do I get the bounding boxes from the resulting mask?

[@Azerty-v8z]

Thanks for this amazing share. Is there any possibility SAM output the label associated with predicted mask in order to know the name of the instance segmented using SAM please? Thanks in advance

[@ariharasudhanmuthusami2272]

Is it possible to use text prompts for fine tuning?

[@user-rr9jd6jd8l]

Hello, Thank you for giving video to help how to fine tune! I have a error that "ValueError: Unsupported number of image dimensions: 2" In here example = train_dataset[0] for k,v in example.items(): print(k,v.shape) How can i solve it?

[@timanb2491]

if we already have prompt(mask) for test image as an input, why we use SAM to get the mask ? I mean - we already have an answer, how using SAM will help us?

[@user-tp6xo5ew4k]

Thank you for the video, your videos are always helpful! I'm facing this error and can't find a solution. In block 16, when accessing 'train_dataset[0]', I encounter the error: 'ValueError: Unsupported number of image dimensions: 2'. Skipping the block doesn't help as the same error occurs during training. I've searched online but couldn't find anything useful. I'm using Google Colab and these library versions: transformers 4.39.0.dev0, torch 2.1.0+cu121, datasets 2.18.0. I would greatly appreciate it if you could help me solve this problem. Thanks in advance.

[@adikrish6926]

I'm having the same issue, how did you solve it?

[@user-rr9jd6jd8l]

@@adikrish6926 Same here! anybody solved it?

[@user-tp6xo5ew4k]

@@adikrish6926 I haven't figured it out yet, have you?

[@adikrish6926]

Yes I figured it out. The solution was to simply convert the grayscale images to RGB images by reshaping their arrays. The masks still need to stay as grey scale though.

[@AakashGoyal25]

def __getitem__(self, idx): item = self.dataset[idx] image = item["image"] image = np.array(image) # Check if the image is grayscale and convert it to RGB if image.ndim == 2: # Image is grayscale image = np.expand_dims(image, axis=-1) # Expand dimensions to (H, W, 1) image = np.repeat(image, 3, axis=2) # Repeat the grayscale values across the new channel dimension ground_truth_mask = np.array(item["label"]) # Get bounding box prompt prompt = get_bounding_box(ground_truth_mask) # Prepare image and prompt for the model inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt") # Remove batch dimension which the processor adds by default inputs = {k: v.squeeze(0) for k, v in inputs.items()} # Add ground truth segmentation inputs["ground_truth_mask"] = ground_truth_mask return inputs Here is the code for it. This works for me. I hope it will work for you as well.

[@gytisbernotas1610]

Hi! This was great - thank you very much for the tutorial! I was also trying to extend your work and work with the RGB rather than single-channel ones. I adjusted the code to deal with the RG images; however, I don't think I have it right for the loss calculations since I am getting a huuuge negative loss value. I was wondering if you have attempted to work with the RGB images as well?

[@user-vz3df2sp4y]

Hello. I also need to work with RGB data. Could you please your modified code with me?

[@supriyoghosh2003]

Is there any progress on it?

[@FelixWei-rn4bt]

Have you already figured out why the loss function has such a high negative value? I have the same problem

[@johanhaggle7949]

How can you know if you overtrain?

[@pyroswolf8203]

Hi, Thanks for the video, is there a option that we can add point prompts ?

[@ortiznicola8022]

hello, I'm trying to do that right now. Please tell me if you were able to do it

[@mohansantokhi3434]

Where in the notebook segment-anything repo is used.

[@ericbader7998]

Thanks for sharing the video! At 1:44, you mention SAM is designed to take text prompt describing what should be segmented. I am not sure that is the case, can you explain how?

[@kanishkbashyam5293]

Its called langsam. You can find it by search for segment-geospatial. I think it works by using a combination of object-detection and segmentation. The object detection is done with Grounding Dino, which return a bunch of bounding boxes. The object inside these bounding boxes are then segmented using SAM.

[@user-lz2ww8uu8q]

Great work, but I have some trouble. Instead of the example images you provided, I have used mine which are 200x200. However, I have encountered two problems: - The images have to be in grayscale if they are RGB the program stops working in "batch = next(iter(train_dataloader))" - The images have to be 256x256. If I use my 200x200 grayscale images it crashes when training, more specifically when calculating the loss. It says that the ground truth is 200x200, and the prediction is 256x256. Do you know how I can fix this problem?

[@NicolaRomano]

My guess is you can just zero pad your image and it should work (np.pad makes that very easy)

[@user-lz2ww8uu8q]

@@NicolaRomano Thank you! Could you handle work with RGB images?

[@NicolaRomano]

@@user-lz2ww8uu8q you should definitely be able to, I haven't tried honestly, you'll probably simply need to take into account the different shape of the image (e.g. (3,256,256) instead of (256,256)). But also, it depends what you want to do (e.g. do you need segmenting the three channels together or separately?)

[@BuseYaren]

Thanks a lot for the informative video! Do you have any videos applying MedSAM3D?

[@DigitalSreeni]

Not yet!

[@shamukshi]

can you do freelancing ? "solar panel counting from UAV images using SAM"

[@urzdvd]

Great tutorial as always Sreeni, thank you, There is a project called medical SAM, that is already custom training with thousands of medical images, to check it out. In social media you have mentioned a tutorial to pass from binary image to polygon masks. Is there any resource that I can base myself on to do this process?

[@DigitalSreeni]

Converting annotations will be my focus for the next video - hoping to release it on Sep 20th. I need to collect my code from different projects and put it together into a single video tutorial. Please stay tuned :)

[@urzdvd]

@@DigitalSreeni thank you Sreeni, I'll stay tuned.

[@tektronix475]

hi sreeni n ppl! does anyone know about any computer vision ML online forum, to post related questions?. Thx!

[@billlee2641]

May I know where is the 12 images tif? the website only gives us two sets of tif, each have 165 images

[@danieleneh3193]

Good day Sir please is it possible to us the SamautomaticMaskgenerator with fine tuned model please how can we generate the mask in the same way SamautomaticMaskgenerator works.

[@mohammed-yassinebarnicha]

can someone please explain to me how can i use this model in the same context but with multiple classes i'm trying to train a sam mode on the fickr material dataset so that it detects materials composing objects

[@timanb2491]

how to unpatch the images?

[@sanjanakala5723]

Hi, How can we train SAM with RGB images and masks like dubai aerial segmentation dataset , can you help with some feedbacks?

[@user-vz3df2sp4y]

Hello. I also want to modify the code for RGB images. Did you successfully execute the code?

[@manalkamal1795]

Hi i have used your code in order to fine tune sam in order to segment aerial images , but when i use my finetunedsam.pth it doesn’t even segment the images that it used to segment with no finetuning, what do you think is the problem ? Thank you in advance !!

[@sulaimanmahmoud7120]

Thanks for great video Is the same way can I apply it on multi class

[@DigitalSreeni]

Sorry, I haven't tested this for multi-class.

[@jww1027]

Please post a video on deep image prior.Thanks

[@mmd_punisher]

Hey man, nice job, u e amazing like a what. I have got a problem in 26:00 min in video, in that 'example' i have an error that says, if anyone can help me, i really appreciate that. this is the last part of ERROR: ...raise ValueError(f"Unsupported number of image dimensions: {image.ndim}") ValueError: Unsupported number of image dimensions: 2

[@lee-qk2vk]

i have the same problem... i wish he did this on spyder ide so we could see the variable explorer. i need to see the dimensions of the input images and masks (hope he can give an answer soon)

[@mmd_punisher]

@@lee-qk2vk The data that returns, is a dic that has 2 keys. also we can use '.dataset' whit that, but i don't really know what i gonna do, also in 2 or 3 lines later, we have this bunch of the code : "batch = next(iter(train_dataloader))" also with same error. hope someone help...

[@Theredeemer-wc6ly]

got the same error

[@mmd_punisher]

@@Theredeemer-wc6ly Uh mate

[@Theredeemer-wc6ly]

@@mmd_punisher there was a fix a few comments ahead

[@danieleneh3193]

Please can you make a video on fine tuning for coco.json data set. Is it possible to fine tune the model for multi-class images

[@ManikandanSathiyanarayanan]

Hi sreeni, great video it is very helpful for me. i was trying to fine tune model for my own custom data but it has 3 channels. while preparing Pytorch custom dataset i had error like "ValueError: zero-size array to reduction operation minimum which has no identity". can you help me to sort out this issue?

[@DigitalSreeni]

This error probably refers to one of your training masks being blank. Try to sort your masks so you only use the ones where you have some information, otherwise the tensor would be empty.

[@ManikandanSathiyanarayanan]

Hi sreeni Thanks for your reply. I have trained SAM model for RGB image but prediction result was empty . can you please tell me what could be wrong? @@DigitalSreeni

[@suzystone4270]

I am trying this tutorial on Breast-Ultrasound-Images-Dataset on Kaggle, I get the same error message during creating a DataLoader instance. When I try to convert to mask into np.array to get the ground_truth_seg, np_unique(ground_truth_seg) does not output array([0, 1], dtype=int32). Instead it outputs an array of bunch of numbers and dtype is. uint8 instead.

[@suzystone4270]

@@DigitalSreeni Thank you! Yes I was getting the same error as I mentioned before and it was because of the blank masks. I filtered them and the error went away.

[@user-vz3df2sp4y]

Hello. I also need to work with RGB data. Could you please your modified code with me?

[@johanhaggle7949]

When changing patch_size from 256 to 512 and step size from 256 to 512 I get this error: "Error: AssertionError: ground truth has different shape (torch.Size([2, 1, 512, 512])) from input (torch.Size([2, 1, 256, 256]))" Why is this?

[@carlosjarrin3170]

There is a part in the image processor class of the 'from transformers import SamProcessor' where it calls a function, and it is stated that the default maximum patch size is 256x256. It took a couple of hours to realize, and I hope it will help somebody. I encourage everyone who wants to understand the code to check the code libraries

[@FelixWei-rn4bt]

@@carlosjarrin3170 is there any chance to use a bigger patch size or is fine- tuning SAM only possible with 256x256? Maybe by using another image processor?

[@Fourest-ys1wi]

@@FelixWei-rn4bt I tried to scale the predicted_masks. And it worked for me. Try this: predicted_masks = outputs.pred_masks.squeeze(1) gt_shape = (640, 640) # the shape of your patch interpolated_mask = F.interpolate(predicted_masks, gt_shape, mode="bilinear", align_corners=False) predicted_masks = interpolated_mask.float()

[@jacobidoko3924]

How does this model compare to the nnUNetv2 model?

[@md.shafiqulislam5692]

Great Tutorial. can you share your notebook?

[@DigitalSreeni]

github.com/bnsreenu/python_for_microscopists/blob/master/331_fine_tune_SAM_mito.ipynb

[@AnkurDe-nz9in]

Hey there! Great work. I came across this video while researching about Segmentation using Transformers. However, on my dataset I am facing a problem. In the cell train_dataset = SAMDataset(dataset=dataset, processor=processor) example = train_dataset[0] for k,v in example.items(): print(k,v.shape) I am getting an error which says Unsupported number of image dimensions: 2. I am using grayscale images here and have tried expanding the dimension of the images while reading it, only to give the same error. If anyone has any suggestion or is aware of some update I have missed, then please go on ahead and educate me :). Am in dire need of some help. Thanks.

[@shubhsinghal8258]

predicted_masks = outputs.pred_masks.squeeze(1) ground_truth_masks = batch["ground_truth_mask"].float().to(device) loss = seg_loss(predicted_masks, ground_truth_masks.unsqueeze(1)) can you explain the output shapes and why ground_truth masks are unsqueezed?

[@strongwarrior0210]

Kindly run df-gan and hifi-gan code. Your code videos are really helpful please help me in running these codes

[@valenparraful]

Hello DigitalSreeni, thank you for this tutorial. I'm getting an error and it's driving me crazy, because I am running your notebook and the same dataset. Everything runs fine, getting exactly the same results, up to the moment where we check an example from the dataset: example = train_dataset[0] for k,v in example.items(): print(k,v.shape) I am getting the following error (Unsupported number of image dimensions: 2): ValueError Traceback (most recent call last) Cell In[17], line 1 ----> 1 example = train_dataset[0] 2 for k,v in example.items(): 3 print(k,v.shape) Cell In[14], line 24 21 prompt = get_bounding_box(ground_truth_mask) 23 # prepare image and prompt for the model ---> 24 inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt") 26 # remove batch dimension which the processor adds by default 27 inputs = {k:v.squeeze(0) for k,v in inputs.items()} File c:\Users\F72070\Document\FC20-dipnn-sot\env_fc20\Lib\site-packages\transformers\models\sam\processing_sam.py:71, in SamProcessor.__call__(self, images, segmentation_maps, input_points, input_labels, input_boxes, return_tensors, **kwargs) 57 def __call__( 58 self, 59 images=None, (...) 65 **kwargs, 66 ) -> BatchEncoding: 67 """ 68 This method uses [`SamImageProcessor.__call__`] method to prepare image(s) for the model. It also prepares 2D 69 points and bounding boxes for the model if they are provided. 70 """ ... --> 200 raise ValueError(f"Unsupported number of image dimensions: {image.ndim}") 202 if image.shape[first_dim] in num_channels: 203 return ChannelDimension.FIRST ValueError: Unsupported number of image dimensions: 2 Any ideas or suggestions would be very appreciated!

[@davidsolooki3051]

Try this: image = np.expand_dims(image, axis=-1) # Add channel dimension image = np.repeat(image, 3, axis=-1) # Repeat grayscale channel to create 3 channels The SAM Processor expects to get 3 input channels. Adding these above two lines of code to the __getitem__ method in the SAMDataset class should solve this issue. See the full example below ####################################################### from torch.utils.data import Dataset class SAMDataset(Dataset): """ This class is used to create a dataset that serves input images and masks. It takes a dataset and a processor as input and overrides the __len__ and __getitem__ methods of the Dataset class. """ def __init__(self, dataset, processor): self.dataset = dataset self.processor = processor def __len__(self): return len(self.dataset) def __getitem__(self, idx): item = self.dataset[idx] image = item["image"] image = np.expand_dims(image, axis=-1) # Add channel dimension image = np.repeat(image, 3, axis=-1) # Repeat grayscale channel to create 3 channels ground_truth_mask = np.array(item["label"]) # get bounding box prompt prompt = get_bounding_box(ground_truth_mask) # prepare image and prompt for the model inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt") # remove batch dimension which the processor adds by default inputs = {k:v.squeeze(0) for k,v in inputs.items()} # add ground truth segmentation inputs["ground_truth_mask"] = ground_truth_mask return inputs

[@billlee2641]

@@davidsolooki3051 thanks!

[@anbuingoc4495]

Dear, how can i modify to train with input shape (512x512x3). Reply me plz~~~

[@Theredeemer-wc6ly]

x3 means that it is a color image, change it to greyscale so it is 2d. 512 by 512

[@anbuingoc4495]

@@Theredeemer-wc6ly thank you bro for replying me 🙏

[@johanhaggle7949]

What if you have bigger objects than mitochondria so that the patches of 256x256 are to small? In this video (video 206) kzbin.info/www/bejne/gn6cqpypg76Zr9k you say that patches should be at least 4 times bigger than the objects. But what if the object is big and I try to change patch size from 256 to e.g. 512 in your colab script I get this error: "Error: AssertionError: ground truth has different shape (torch.Size([2, 1, 512, 512])) from input (torch.Size([2, 1, 256, 256]))"

[@cXedis]

darkmode please....... for the love of all that is holy.....

[@Jay-kb7if]

what's up with tffs dude.

[@yi9itc4n]

this shi complicated af

[@thiccMaleChicken]

can we fine-tune segment anything model other than base one?