Skip to content

Commit

Permalink
Add SAM decoder & output masks as png (#418)
Browse files Browse the repository at this point in the history
* Add loading of decoder layers in Model

* Multiply by hypernet_layer_cnt for ctx_size on model load

* Add decoder layers to py conversion script

* Fix wrong and reversed tensor sizes for decoder

* Add decoder transformer implementation

* Add decoder hypernet and iou prediction mlps

* Add transpose convolution operation and unit test

* Finish mask decoder and write the decoder output in the model state

* Output masks to png after removing padding and upsampling to original size

- Also filter based on the iou treshold
- Additionally filtering based on the stability score and crop boxes
should be done

* Add stb image write in order to output masks from SAM

* Add transpose convolution 2d name and symbol to ggml ops static arrays

* Comment out debug print in transpose convolution test to fix compilation

ggml-ci
  • Loading branch information
Yavor Ivanov committed Aug 17, 2023
1 parent ee26ea7 commit 03b4917
Show file tree
Hide file tree
Showing 7 changed files with 3,144 additions and 151 deletions.
9 changes: 3 additions & 6 deletions examples/sam/convert-pth-to-ggml.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,6 @@
name = k
shape = v.shape

# TODO: export only the Encoder -- after it works we will export the other stuff
if name[:13] != "image_encoder" and \
name[:14] != "prompt_encoder":
continue

if name[:19] == "prompt_encoder.mask":
continue

Expand Down Expand Up @@ -104,7 +99,9 @@
ftype_cur = 1
if ftype == 0 or n_dims == 1 or \
name == "image_encoder.pos_embed" or \
name.startswith("prompt_encoder"):
name.startswith("prompt_encoder") or \
name.startswith("mask_decoder.iou_token") or \
name.startswith("mask_decoder.mask_tokens"):
print(" Converting to float32")
data = data.astype(np.float32)
ftype_cur = 0
Expand Down
1,174 changes: 1,031 additions & 143 deletions examples/sam/main.cpp

Large diffs are not rendered by default.

Loading

0 comments on commit 03b4917

Please sign in to comment.