get_objects.py<\/code> file and get started.<\/p>\n# USAGE: python get_objects.py\n\n# Import the necessary packages\nimport os\nimport pickle\n\nimport cv2\nimport matplotlib.pyplot as plt\nfrom segment_anything import SamAutomaticMaskGenerator, sam_model_registry\n\nfrom pyimagesearch import config, utils\n\n\ndef generate_object_masks(image):\n # Initialize SAM model and generate masks\n print(\"[INFO] Loading SAM model...\")\n sam = sam_model_registry[config.MODEL_TYPE](checkpoint=config.SAM_CHECKPOINT_PATH)\n mask_generator = SamAutomaticMaskGenerator(sam)\n masks = mask_generator.generate(image)\n\n return masks\n\n\ndef save_object_crops(masks, mask_ids, labels):\n # Crop objects from image using mask IDs and save\n obj_crops = []\n for id in mask_ids:\n box = masks[id][\"bbox\"]\n x0, y0, w, h = box\n crop = image[y0 : y0 + h, x0 : x0 + w]\n obj_crops.append(crop)\n\n # Save object crops and corresponding labels\n obj_dict = {\"crops\": obj_crops, \"labels\": labels}\n with open(os.path.join(config.OUT_PATH, \"objects.pkl\"), \"wb\") as fp:\n pickle.dump(obj_dict, fp)<\/pre>\nWe start by importing the os<\/code> module, and the pickle<\/code> package (Lines 4 and 5<\/strong>). Next we import the OpenCV and Matplotlib libraries, as shown on Lines 7 and 8<\/strong>, respectively. On Line 9<\/strong>, we import necessary modules that will allow us to use SAM for making predictions in our tutorial. Specifically, we get the SamAutomaticMaskGenerator<\/code> and sam_model_registry<\/code> modules from segment_anything<\/code>, as shown on Line 9<\/strong>.<\/p>\nFinally we import the config<\/code> file, which stores the initial configurations of our parameters (Line 11<\/strong>) and the utils.py<\/code> file, which implements the helper functions to visualize the segmentation masks from SAM (Line 11<\/strong>).<\/p>\nNow that we have imported the necessary modules and packages, it is time to implement our get_object_masks<\/code> function (Lines 14-21<\/strong>), which takes the input image and uses the pre-trained SAM to predict and automatically generate masks for the entire image.<\/p>\nWe first initialize SAM using the sam_model_registry<\/code> functionality with the type of model (i.e., config.MODEL_TYPE<\/code>), which we take as vit_h<\/code> for this tutorial. Further, we provide the path where we stored the SAM checkpoint (i.e., config.SAM_CHECKPOINT_PATH<\/code>), as shown on Line 17<\/strong>. <\/p>\nOn Line 18<\/strong>, we get the mask_generator<\/code> using the SamAutomaticMaskGenerator<\/code> module, allowing us to automatically extract masks for the different objects in our input image. <\/p>\nFinally, on Line 19<\/strong>, we use the generator functionality of the mask_generator<\/code> object that we created to get the masks for the entire image and return the predicted masks<\/code> on Line 21<\/strong>.<\/p>\nNow that we have completed the definition of our get_object_masks<\/code> function, discuss the save_object_crops<\/code> function which will allow us to crop out everyday objects detected by SAM from our input image and save them with their corresponding labels. This will further help us perform the downstream tasks we want with the different objects in our image as we will see later.<\/p>\nWe first initialize an empty obj_crops<\/code> list (Line 26<\/strong>), where we will store our object crops. Then, for each ID in our mask_ids<\/code> list, we get the bounding box predicted by SAM (Line 28<\/strong>) and get the corresponding x0<\/code>, y0<\/code> coordinates and width and height (Line 29<\/strong>).<\/p>\nNext, we use the coordinates to get the part or crop of the image with the particular object (Line 30<\/strong>) and add it to our obj_crops<\/code> list (Line 31<\/strong>).<\/p>\nFinally, we create a dictionary named obj_dict<\/code> and store the obj_crops<\/code> list and the corresponding labels, as shown on Line 34<\/strong>. Then, we save this dictionary with the filename objects.pkl<\/code> in the config.OUT_PATH<\/code> folder (Line 35<\/strong>).<\/p>\nif __name__ == \"__main__\":\n # Load input image and convert to RGB\n print(\"[INFO] Loading image...\")\n image = cv2.imread(config.IMG_PATH[1])\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n # Generate segmentation masks\n print(\"[INFO] Generating masks from SAM...\")\n all_masks = generate_object_masks(image)\n\n # Create output directory if it doesn't exist\n if not os.path.exists(config.OUT_PATH):\n os.makedirs(config.OUT_PATH)\n\n # Plot and save predicted image\n prediction_path = os.path.join(config.OUT_PATH, config.OUT_PRED_PATH)\n plt.figure(figsize=(8, 8))\n plt.imshow(image)\n utils.show_all_masks(all_masks)\n plt.savefig(prediction_path)\n\n # Plot individual object masks\n for i, mask_info in enumerate(all_masks):\n plt.figure(figsize=(6, 6))\n plt.imshow(image)\n utils.show_mask(mask_info[\"segmentation\"], plt.gca())\n plt.axis(\"off\")\n plt.savefig(os.path.join(config.OUT_PATH, f\"mask_{i}.png\"))\n\n # TODO: Place mask and labels in the config\n mask_ids = [0, 2, 3, 7, 8, 11, 12, 20, 22, 23, 26, 27, 34, 35, 45, 66, 67, 78]\n labels = [\n \"painting\",\n \"door\",\n \"painting\",\n \"carpet\",\n \"vase\",\n \"plant\",\n \"pillow\",\n \"blanket\",\n \"pillow\",\n \"pillow\",\n \"blanket\",\n \"jug\",\n \"blanket\",\n \"dog\",\n \"plant\",\n \"plant\",\n \"plant\",\n \"pillow\",\n ]\n\n # Save object crops and corresponding labels\n print(\"[INFO] Cropping and saving objects...\")\n save_object_crops(all_masks, mask_ids, labels)<\/pre>\nNext, let us go ahead and implement the main function to see our model in action.<\/p>\n
We start by loading our image using OpenCV from the path defined by config.IMG_PATH[1]<\/code> and converting it from BGR to RGB color space using the cv2.cvtColor<\/code> function, as shown on Lines 42 and 43<\/strong>. <\/p>\nNext, on Line 47<\/strong>, we use the get_object_masks<\/code> function that we defined above to extract the segmentation masks from our input image and store it as all_masks<\/code>. Notice that the all_masks<\/code> output from SAM is a dictionary with the following keys: [‘segmentation’, ‘area’, ‘bbox’, ‘predicted_iou’, ‘point_coords’, ‘stability_score’, ‘crop_box’].<\/p>\nNote that we will use the segmentation masks (i.e., ‘segmentation’ key) and the bounding box (i.e., ‘bbox’ key) for cropping individual objects.<\/p>\n
Now that we have the SAM predictions, let us prepare to visualize them on our input image and save the final visualizations.<\/p>\n
We first check if the folder where the output predictions will be stored (i.e., config.OUT_PATH<\/code>) already exists, and if not, we create it (Lines 50 and 51<\/strong>). Furthermore, we define the prediction_path<\/code> on Line 54<\/strong>, which indicates the location where our segmentation visualization will be stored.<\/p>\nOn Lines 55 and 56<\/strong>, we initialize a matplotlib figure and visualize the input image.<\/p>\nNext, we visualize all the predicted segmentation masks from SAM using the show_all_masks<\/code> function and save the visualization using the plt.savefig<\/code> at the prediction_path<\/code>, as shown on Lines 57 and 58<\/strong>.<\/p>\nFigure 2<\/strong> shows the visualization of all segmentation masks predicted by SAM for our input image.<\/p>\n\n
Figure 2:<\/strong> Input Image (left<\/em>) and all mask prediction from SAM (right<\/em>) (source: image by the author).<\/figcaption><\/figure>\n<\/div>\nNow, let us go ahead and visualize each segmentation mask and the corresponding bounding box predicted by SAM to get a sense of the prominent objects that were segmented out by SAM.<\/p>\n
We start by iterating over the all_masks<\/code> output. For each segmentation prediction, we first plot the input image (Lines 62 and 63<\/strong>) and then visualize the segmentation mask (i.e., mask_info['segmentation']<\/code>) using the show_mask<\/code> function (Lines 64 and 65<\/strong>).<\/p>\nFinally, we save the visualization using the plt.savefig<\/code> in the config.OUT_PATH<\/code> folder, as shown on Line 66<\/strong>.<\/p>\nNow that we have visualized each prediction and their corresponding segmentation masks and bounding boxes, let us take a list of predicted mask IDs from SAM with prominent everyday objects and the corresponding object names or labels (Lin<\/strong>e 69<\/strong>).<\/p>\nNote that on Lin<\/strong>e 69-89<\/strong>, we manually curate a few mask IDs (which we visualized) with prominent everyday objects (e.g., paintings, plants, vases, pillows, etc.) and assign them corresponding labels with the help of the labels list.<\/p>\nFinally we use the save_object_crops<\/code> function that we defined above to save the object crops.<\/p>\nWe will use these objects extracted from our input image for our downstream tasks.<\/p>\n
Figure 3<\/strong> shows a few examples of cropped objects and their corresponding labels. <\/p>\n
![\"\"](\"https:\/\/b2633864.smushcdn.com\/2633864\/wp-content\/uploads\/2023\/09\/sam-2-featured.png?lossy=2&strip=1&webp=1\")
<\/a><\/figure>\n<\/div>\n![\"\"](\"https:\/\/b2633864.smushcdn.com\/2633864\/wp-content\/uploads\/2020\/01\/source-code-icon.png?lossy=2&strip=1&webp=1\"){kind=icon}
<\/p>\n