add sa-1b dataset

Author: Mark-ZhouWX

official/cv/segment-anything/configs/sa1b_box_finetune.yaml

@@ -0,0 +1,76 @@
+#---------------------------------------------
+# Part 1: system basic config setting
+distributed: False
+device: Ascend
+mode: 0  # 0: graph, 1: pynative
+work_root: &work_root ./work_dir/
+log_level: info
+amp_level: O2
+
+# ---------------------------------------------
+# Part2: module setting
+loss_manager:
+#  type: fixed  # dynamic or
+#  scale_sense: 1024
+  loss_scaler:
+    type: dynamic
+  grad_clip: False
+  drop_overflow_update: False
+
+optimizer:
+  type: segment_anything.optim.optimizer.AdamW
+  weight_decay: 1e-4
+  group_param:
+
+  lr_scheduler:
+    type: segment_anything.optim.scheduler.SAMDynamicDecayLR
+    learning_rate: 8e-6
+    warmup_steps: 250
+    decay_steps: [ 60000, 86666 ]
+    decay_factor: 10
+
+
+network:
+  model:
+    type: vit_b
+    checkpoint: ./models/sam_vit_b-35e4849c.ckpt
+    freeze:
+      image_encoder: True
+      prompt_encoder: True
+
+  loss:
+    type: segment_anything.modeling.loss.SAMLoss
+
+
+train_loader:
+  dataset:
+    type: segment_anything.dataset.dataset.SA1BDataset
+    data_dir: ./datasets/sa-1b/
+    transform_pipeline:
+      - type: segment_anything.dataset.transform.ImageResizeAndPad
+        target_size: 1024
+      - type: segment_anything.dataset.transform.ImageNorm
+        hwc2chw: True
+      - type: segment_anything.dataset.transform.LabelPad
+        gt_size: 20
+    output_column: ['image', 'masks', 'boxes', 'valid_boxes']
+
+  model_column: ['image', 'boxes']  # columns for model cell input
+  loss_column:  ['masks', 'valid_boxes']  # columns for loss function input
+
+  shuffle: True
+  batch_size: 1
+  epoch_size: 8
+  drop_remainder: True
+  num_workers: 2
+  max_rowsize: 64  # 24M space for dataloader
+
+
+callback:
+  - type: segment_anything.utils.callbacks.TrainStatusLog
+    loss_item: ['focal_loss', 'dice_loss', 'mse_loss']  # for log
+    interval: 100
+  - type: segment_anything.utils.callbacks.SaveCkpt
+    work_root: *work_root
+    interval: 1  # in epoch
+

official/cv/segment-anything/segment_anything/dataset/dataset.py

@@ -1,3 +1,4 @@
+import json
 import os
 from typing import List
 
@@ -6,6 +7,7 @@
 from mindspore.dataset import GeneratorDataset, BatchDataset
 
 from pycocotools.coco import COCO
+from pycocotools import mask as maskUtils
 
 from segment_anything.dataset.transform import create_transform_pipeline
 from segment_anything.utils import logger
@@ -142,3 +144,98 @@ def __getitem__(self, idx):
             self.output_column = list(data_dict.key())
 
         return tuple(data_dict[k] for k in self.output_column)
+
+
+@DATASET_REGISTRY.registry_module()
+class SA1BDataset:
+
+    def __init__(self,
+                 data_dir,
+                 transform_pipeline,
+                 output_column: List[str] = None,
+                 **kwargs,
+                 ):
+        self.data_dir = data_dir
+        self.output_column = output_column
+        self.transform_pipeline = create_transform_pipeline(transform_pipeline)
+        assert os.path.exists(data_dir), f'SA-1B dataset root not exists at {data_dir}'
+        parts = sorted(os.listdir(data_dir))  # there are about 11K jpgs in each part
+
+        image_paths = []
+        anno_paths = []
+        for p in parts:
+            part_dir = os.path.join(data_dir, p)
+            all_files = [os.path.join(part_dir, f) for f in sorted(os.listdir(part_dir))]
+            image_paths += list(filter(lambda f: f.endswith('.jpg'), all_files))
+            anno_paths += list(filter(lambda f: f.endswith('.json'), all_files))
+            assert len(image_paths) == len(anno_paths)
+
+        self.image_paths = image_paths
+        self.anno_paths = anno_paths
+
+        logger.info(f'got {len(parts)} parts of SA-1B dateset, total size: {len(self.image_paths)}')
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        """
+        Below is an example describing the format of sa-1b dataset
+            annotations:
+                - area: 1418,
+                  bbox: [1134.0, 119.0, 30.0, 58.0]
+                  crop_box: [622.0, 0.0, 567.0, 707.0]  # the sub part of image where the mask is generated, see automatic_mask_generator.py for details
+                  point_coords: [1153.5625, 132.5625]
+                  predicted_iou: 0.8891242146492
+                  segmentation:
+                    counts: "`]YW23SP2`0D;F:F8H3M1O100O1O1O10000O1100O001O001O1O010O1O001O1^OYQN\\Oin16iQNEXn14X1FbiQe0"
+                    size: [2060, 1500]
+                  stability_score: 0.960608184337616
+                - area:
+                  box:
+                  xxx:
+            ]
+            image:
+                file_name: "sa_1.jpg"
+                height: 2060
+                image_id: 1
+                width: 1500
+        Returns:
+            a tuple of transformed input items
+
+        """
+        anno_path = self.anno_paths[idx]
+        image_path = self.image_paths[idx]
+
+        assert os.path.exists(image_path), f'image file not found at {image_path}'
+        assert os.path.exists(anno_path), f'anno file not found at {anno_path}'
+
+        image = cv2.imread(image_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        with open(anno_path, 'r') as f:
+            json_data = json.load(f)
+        anno_list = json_data['annotations']
+
+        boxes = []
+        masks = []
+        for anno in anno_list:
+            x, y, w, h = anno['bbox']
+            mask = maskUtils.decode(anno['segmentation'])  # uint8
+
+            # filter small mask
+            image_h, image_w = anno['segmentation']['size']
+            if w / image_w < 0.1 and h / image_h < 0.1:
+                continue
+
+            boxes.append([x, y, x + w, y + h])
+            masks.append(mask)
+
+        # letter box
+        data_dict = dict(image=image, masks=masks, boxes=np.array(boxes, np.float32))
+        data_dict = self.transform_pipeline(data_dict)
+
+        if self.output_column is None:
+            self.output_column = list(data_dict.key())
+
+        return tuple(data_dict[k] for k in self.output_column)

official/cv/segment-anything/segment_anything/dataset/transform.py

@@ -77,7 +77,7 @@ def __call__(self, result_dict):
         if False:  # show image and mask for debug
             import matplotlib.pyplot as plt
             plt.imshow(result_dict['image'])  # raw image
-            from use_sam_with_promts import show_box, show_mask
+            from segment_anything.utils.visualize import show_box, show_mask
             show_box(result_dict['boxes'][0], plt.gca())
             show_mask(result_dict['masks'][0], plt.gca())
             plt.show()