article SAM2 on Databricks in Technical Blog

article SAM2 on Databricks in Technical Blog https://community.databricks.com/t5/technical-blog/sam2-on-databricks/ba-p/92497 <P><SPAN>In this post we’ll walk through getting started with Meta’s latest Segment-Anything-Model 2 (SAM2) on Databricks. We’ll cover experimentation with SAM2 in a Databricks Notebook, expand on the default examples by using custom notebook cells for visualization.</SPAN></P> <P> </P> <P>Want to skip straight to the code? Check out the complete <A href="https://github.com/databricks-solutions/devrel-examples/blob/main/notebooks/sam2-on-databricks.ipynb" target="_self">notebook</A> to get started.</P> <P><STRONG>What is SAM2?</STRONG></P> <P><SPAN>SAM 2 is a unified model for segmenting objects across both images and videos. It supports point, box, and mask coordinates to select an object on any image or frame of video, and shows significant improvements over SAM in performance, accuracy, and manual intervention in both video and image tasks. More details are available in the original research paper published </SPAN><A href="https://ai.meta.com/research/publications/sam-2-segment-anything-in-images-and-videos/" target="_blank" rel="noopener"><SPAN>here</SPAN></A><SPAN>.</SPAN></P> <P><STRONG>Why do segmentation?</STRONG></P> <P><SPAN>Segmentation has many applications across various industries. Some examples include object detection in autonomous vehicles, medical image analysis for disease diagnosis, and content moderation in social media platforms. </SPAN><SPAN><BR /><BR /></SPAN></P> <P><SPAN>While many hosted services offer segmentation capabilities as part of their machine learning offerings, self-deploying segmentation models like SAM2 can provide significant advantages: greater control over the segmentation process, the ability to customize the model for specific use cases and potential cost savings for high-volume applications. Self-deployment also enables full self governance, with better data privacy and security, as sensitive images or videos don't need to be sent to third-party services for processing.</SPAN></P> <P><STRONG>SAM2 on Databricks</STRONG></P> <P><SPAN>To get started we’ll </SPAN><A href="https://docs.databricks.com/en/compute/gpu.html" target="_blank" rel="noopener"><SPAN>need a running cluster with GPU support</SPAN></A><SPAN>. We’ve chosen to use the </SPAN><STRONG>g2dn.xlarge [T4]</STRONG><SPAN> with Databricks Runtime Version </SPAN><STRONG>15.4 LTS ML</STRONG><SPAN>, but you can choose a beefier setup for bigger workloads.</SPAN></P> <P><SPAN>Next, checkout </SPAN><A href="https://github.com/databricks-solutions/devrel-examples/blob/main/notebooks/sam2-on-databricks.ipynb" target="_blank" rel="noopener"><SPAN>this notebook </SPAN></A><SPAN>from the Databricks Dev Rel samples repository. You can run through the notebook directly in Databricks but we’ll highlight some of the key sections below.</SPAN></P> <P><STRONG>Setup</STRONG></P> <P><SPAN>The easiest way to to run SAM2 on Databricks is to create a new Databricks Git Folder configured with the SAM2 Github (<A href="https://github.com/facebookresearch/segment-anything-2" target="_blank" rel="noopener">https://github.com/facebookresearch/segment-anything-2</A>) and import the notebook above into the existing notebooks folder within the created Git folder </SPAN><SPAN>segment-anything-2/notebooks</SPAN><SPAN>. Paths in the notebook assume the correct location of the notebook within the repository.</SPAN></P> <P><SPAN>We’ll install the SAM2 library directly from Github and use the </SPAN><SPAN>download_ckpts.sh</SPAN><SPAN> to download the model binaries. By default the script will download all sizes of the model, but you can modify the file directly to download a specific size. The sizes available are tiny, small, base_plus, and large.</SPAN></P> <P> </P> <LI-CODE lang="python"># install the sam2 pip package directly from the Databricks Git folder %sh pip install ../../segment-anything-2 # make sure to build extensions %sh cd ../ && python setup.py build_ext --inplace # download the model binaries %sh cd ../checkpoints && ./download_ckpts.sh</LI-CODE> <P> </P> <P><SPAN>You can use any video you like, but in this post we’ll use the Day 1 Keynote from the 2024 DATA & AI Summit. We’ll need to preprocess and break this video up into individual frames and upload them directly to Databricks. We’ll use </SPAN><A href="https://www.ffmpeg.org/" target="_blank" rel="noopener"><SPAN>FFmpeg</SPAN></A><SPAN> and a </SPAN><A href="https://docs.databricks.com/en/sql/language-manual/sql-ref-volumes.html" target="_blank" rel="noopener"><SPAN>Databricks Volume</SPAN></A><SPAN> in this example but it’s not a requirement.</SPAN></P> <P> </P> <LI-CODE lang="python"># use yt-dlp to get the keynote video from youtube %pip install yt-dlp %sh yt-dlp -o ./videos/keynote/keynote.mp4 -f "bestvideo[height<=480]" -u "username" -p "password" "https://www.youtube.com/watch?v=-6dt7eJ3cMs" # use ffmpeg to split the video into frames and place them in a Volume %sh ffmpeg -ss 00:00:15 -i ./videos/keynote/keynote.mp4 -t 00:00:10 -q:v 2 -start_number 0 /Volumes/sam/default/frames/'%05d.jpg'</LI-CODE> <P> </P> <P><STRONG>Segmentation</STRONG></P> <P><SPAN>You can use a combination of Python I/O and matplotlib to visualize the frames directly in the notebook. This is directly from the examples in Meta’s SAM2 examples</SPAN></P> <P> </P> <LI-CODE lang="python"># `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg` video_dir = "/Volumes/sam/default/frames" # scan all the JPEG frame names in this directory frame_names = [ p for p in os.listdir(video_dir) if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] ] frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) # take a look the first video frame frame_idx = 0 plt.figure(figsize=(12, 8)) plt.title(f"frame {frame_idx}") plt.imshow(Image.open(os.path.join(video_dir, frame_names[frame_idx])))</LI-CODE> <P> </P> <P><SPAN><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Screenshot 2024-10-01 at 2.14.08 PM.png" style="width: 400px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/11632i5AF1C6E4FCE5D3EA/image-size/medium?v=v2&px=400" role="button" title="Screenshot 2024-10-01 at 2.14.08 PM.png" alt="Screenshot 2024-10-01 at 2.14.08 PM.png" /></span></SPAN></P> <P><SPAN>You need to initialize the predictor ahead of time. This will process all the images in the Volume that were created with ffmpeg:</SPAN></P> <P> </P> <LI-CODE lang="python">inference_state = redictor.init_state(video_path=video_dir)</LI-CODE> <P> </P> <P><SPAN>You can use just a single point to start segmenting. In the image above matplotlib gives us a coordinate system so we see that the head is at roughly (350,100).</SPAN></P> <P><SPAN>To improve segmentation results you can register as many points as you want or even the coordinates of a larger polygon.</SPAN></P> <P> </P> <LI-CODE lang="python">points = np.array([[350, 100]], dtype=np.float32) ... _, out_obj_ids, out_mask_logits = predictor.add_new_points( inference_state=inference_state, points=points, ... )</LI-CODE> <P> </P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Screenshot 2024-10-01 at 2.15.38 PM.png" style="width: 400px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/11633iE248E3262B1EC552/image-size/medium?v=v2&px=400" role="button" title="Screenshot 2024-10-01 at 2.15.38 PM.png" alt="Screenshot 2024-10-01 at 2.15.38 PM.png" /></span></P> <P><SPAN>After propagating the original detection through the remaining video we can see that it successfully follows the head.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="Screenshot 2024-10-01 at 2.17.09 PM.png" style="width: 335px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/11634i279E7D732868193B/image-size/medium?v=v2&px=400" role="button" title="Screenshot 2024-10-01 at 2.17.09 PM.png" alt="Screenshot 2024-10-01 at 2.17.09 PM.png" /></span></P> <P><SPAN>Finally, we can use both IPython Widgets and Databricks Notebook </SPAN><SPAN>displayHTML</SPAN><SPAN> feature to customize views to go through our frames. In the linked notebook we demonstrate a built-in JS snippet to display the coordinates of our mouse as we move it around the frame, which assists in finding the coordinates of objects.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="mouseover.gif" style="width: 400px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/11635iF8FAF32D63E26FAB/image-size/medium?v=v2&px=400" role="button" title="mouseover.gif" alt="mouseover.gif" /></span></P> <P><STRONG>Conclusion</STRONG></P> <P><SPAN>In this blog we’ve introduced the SAM2 on Databricks example to get started with segmenting using the features of Databricks Notebooks. You can take this further by combining object detection models to find the original coordinates of objects for a truly hands off segmentation capability. You can find the latest on the </SPAN><A href="https://huggingface.co/models?pipeline_tag=object-detection&sort=trending" target="_blank" rel="noopener"><SPAN>Huggingface Models</SPAN></A><SPAN> directory.</SPAN></P> Tue, 01 Oct 2024 21:23:30 GMT NickKarpov 2024-10-01T21:23:30Z SAM2 on Databricks https://community.databricks.com/t5/technical-blog/sam2-on-databricks/ba-p/92497 <P><SPAN>Get started with Meta’s latest Segment-Anything-Model 2 (SAM2) using Databricks Notebooks.</SPAN></P> Tue, 01 Oct 2024 21:23:30 GMT https://community.databricks.com/t5/technical-blog/sam2-on-databricks/ba-p/92497 NickKarpov 2024-10-01T21:23:30Z