Gripper segment custom models #5
Conversation
…sim into gripper-segment
…sim into gripper-segment
.gitignore
Outdated
| @@ -5,3 +5,4 @@ checkpoints | |||
| __pycache__ | |||
| outputs | |||
| .vscode | |||
| ~/ | |||
There was a problem hiding this comment.
nope, forgot to remove that after adding it in from a local git annoyance
| We provide fine tuned networks for SAM2 and GroundingDINO for the segmentation and annotation | ||
| of the gripper used in our provided dataset which can be downloaded from [here](https://mitprod-my.sharepoint.com/personal/nepfaff_mit_edu/_layouts/15/onedrive.aspx?id=%2Fpersonal%2Fnepfaff%5Fmit%5Fedu%2FDocuments%2Fscalable%5Freal2sim%5Fmodel%5Fweights&ga=1). | ||
|
|
||
| Please put the downloaded checkpoint files in the `./checkpoints` directory. |
There was a problem hiding this comment.
:nit: That directory doesn't currently exist. It might make sense to create it and put a .gitignore file inside that includes all files inside this directory apart from the directory itself (you can use ! to exclude the dir).
| To fine tune your own segmentation model for your gripper, see [these instructions](https://github.com/facebookresearch/sam2/blob/main/training/README.md) for training from the | ||
| SAM2 Official Github. | ||
|
|
||
| An example of segmentation failure on the gripper with default models: \ |
There was a problem hiding this comment.
Could you elaborate a bit here for why it can fail without fine-tuning? It could be worth mentioning that our particular gripper seems to be out of distribution for SAM2, and thus it looses track of it for long videos. This can also be solved with re-prompting it after failure, segmenting the video in parts but fine-tuning removes this extra step.
README.md
Outdated
| SAM2 Official Github. | ||
|
|
||
| An example of segmentation failure on the gripper with default models: \ | ||
| <img src="assets/mask_sam2_default.png" width="200"> |
There was a problem hiding this comment.
:nit: Maybe "_failure" is more appropriate than "_default".
| @@ -108,6 +108,30 @@ automatic annotations using DINO work well in many cases but can struggle with t | |||
| gripper masks. All downstream object tracking and reconstruction results are sensitive | |||
| to the segmentation quality and thus spending a bit of effort here might be worthwhile. | |||
|
|
|||
| ##### Gripper masking with fine tuned models | |||
There was a problem hiding this comment.
I think it would make sense to also make this an option for the run_asset_generation.py pipeline. Maybe add a store_true argument to enable it.
In that case, we have to think about where to put the installation instructions. We could keep them here but refer to them here when explaining that argument.
| @@ -112,7 +117,22 @@ def segment_moving_obj_data( | |||
| "Text prompt must be provided if GUI frames are not specified." | |||
| ) | |||
|
|
|||
| sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt" | |||
| if txt_prompt == "gripper": | |||
There was a problem hiding this comment.
I would again check if the install was successful by adding MMDET_AVAILABLE to the if statement. You could print another warning here that mentions that we are falling back to the default model (if `txt_prompt == "gripper" frame).
| sam2_checkpoint = gripper_sam2_path | ||
|
|
||
| if not os.path.exists(sam2_checkpoint): | ||
| logging.info( |
| ) | ||
| if txt_prompt == "gripper": | ||
| # Use mmdetection api for gripper | ||
| with autocast(enabled=False): |
There was a problem hiding this comment.
Do we need autocost if we set enabled=False?
There was a problem hiding this comment.
not sure why this caused an error if I didn't have this here
There was a problem hiding this comment.
In that case, I'd add a comment that mentions that it is needed to avoid errors
| @@ -653,6 +712,16 @@ def get_dino_boxes(text, frame_idx): | |||
| except Exception as e: | |||
| logging.error(f"Error deleting {f}: {e}") | |||
|
|
|||
| # save black masks for gripper not found frames | |||
There was a problem hiding this comment.
:nit: Starting with lower case breaks the file convention
There was a problem hiding this comment.
Also, adding a bit more detail to that comment might be helpful here. Am I understanding it correctly that you are writing black masks until the gripper appears in the frame?
| input_boxes, confidences, class_names = get_dino_boxes( | ||
| txt_prompt, txt_prompt_index | ||
| ) | ||
| while True: |
There was a problem hiding this comment.
:nit: Could you add a comment for what this is doing?
Add support for using custom SAM2 or Grounding DINO models for gripper segmentation
Failed segmentation using default Grounding DINO + SAM2 models:
Successful segmentation using fine tuned Grounding DINO + SAM2 models:
