Spaces:
Runtime error
Runtime error
| # coding=utf-8 | |
| # Copyright 2024 The Google Research Authors. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """The inference pipeline for the CaR model.""" | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| # pylint: disable=g-importing-member | |
| # pylint: disable=g-bad-import-order | |
| from modeling.post_process.post_process import generate_masks_from_sam | |
| from modeling.post_process.post_process import match_masks | |
| from utils.utils import process_sentence | |
| from utils.metrics import IoU | |
| IMAGE_WIDTH = 512 | |
| IMAGE_HEIGHT = 512 | |
| def get_sam_masks( | |
| config, image_path, masks, matching_thresh=0.9, img_sam=None, pipeline=None | |
| ): | |
| """Generate SAM masks.""" | |
| print("generating sam masks online") | |
| mask_tensor, mask_list = generate_masks_from_sam( | |
| image_path, | |
| save_path="./", | |
| pipeline=pipeline, | |
| img_sam=img_sam, | |
| visualize=False, | |
| ) | |
| mask_tensor = mask_tensor.to(masks.device) | |
| # only conduct sam on masks that is not all zero | |
| attn_map, mask_ids = [], [] | |
| for mask_id, mask in enumerate(masks): | |
| if torch.sum(mask) > 0: | |
| attn_map.append(mask.unsqueeze(0)) | |
| mask_ids.append(mask_id) | |
| matched_masks = [ | |
| match_masks( | |
| mask_tensor, | |
| attn, | |
| mask_list, | |
| iom_thres=config.car.iom_thres, | |
| min_pred_threshold=config.sam.min_pred_threshold, | |
| ) | |
| for attn in attn_map | |
| ] | |
| for matched_mask, mask_id in zip(matched_masks, mask_ids): | |
| sam_masks = np.array([item["segmentation"] for item in matched_mask]) | |
| sam_mask = np.any(sam_masks, axis=0) | |
| cur_mask = masks[mask_id] | |
| iou = IoU(torch.from_numpy(sam_mask).to(cur_mask.device), cur_mask) | |
| if iou > matching_thresh: | |
| masks[mask_id] = torch.from_numpy(sam_mask).to(masks.device) | |
| return masks | |
| def inference_car(cfg, car_model, image_path, sentences, sam_pipeline=None): | |
| sentences = [process_sentence(sen, cfg.test.ds_name) for sen in sentences] | |
| img = Image.open(image_path).convert("RGB") | |
| if cfg.test.use_pseudo: | |
| masks, scores = car_model(img, sentences) | |
| return masks, scores | |
| masks, scores = car_model(img, sentences, cfg.car.num_iteration) | |
| sam_masks = get_sam_masks( | |
| cfg, image_path, masks, cfg.sam.matching_thresh, pipeline=sam_pipeline | |
| ) | |
| return sam_masks, scores | |