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. | |
| """DenseCRF.""" | |
| import numpy as np | |
| from pydensecrf import densecrf as dcrf | |
| from pydensecrf import utils | |
| import torch | |
| import torch.nn.functional as F | |
| class DenseCRF(object): | |
| """DenseCRF class.""" | |
| def __init__(self, iter_max, pos_w, pos_xy_std, bi_w, bi_xy_std, bi_rgb_std): | |
| self.iter_max = iter_max | |
| self.pos_w = pos_w | |
| self.pos_xy_std = pos_xy_std | |
| self.bi_w = bi_w | |
| self.bi_xy_std = bi_xy_std | |
| self.bi_rgb_std = bi_rgb_std | |
| def __call__(self, image, probmap): | |
| c, h, w = probmap.shape | |
| u = utils.unary_from_softmax(probmap) | |
| u = np.ascontiguousarray(u) | |
| image = np.ascontiguousarray(image) | |
| d = dcrf.DenseCRF2D(w, h, c) | |
| d.setUnaryEnergy(u) | |
| d.addPairwiseGaussian(sxy=self.pos_xy_std, compat=self.pos_w) | |
| d.addPairwiseBilateral( | |
| sxy=self.bi_xy_std, | |
| srgb=self.bi_rgb_std, | |
| rgbim=image, | |
| compat=self.bi_w, | |
| ) | |
| q = d.inference(self.iter_max) | |
| q = np.array(q).reshape((c, h, w)) | |
| return q | |
| class PostProcess: | |
| """Post processing with dense CRF.""" | |
| def __init__(self, device): | |
| self.device = device | |
| self.postprocessor = DenseCRF( | |
| iter_max=10, | |
| pos_xy_std=1, | |
| pos_w=3, | |
| bi_xy_std=67, | |
| bi_rgb_std=3, | |
| bi_w=4, | |
| ) | |
| def apply_crf(self, image, cams, bg_factor=1.0): | |
| """Apply dense CRF.""" | |
| bg_score = np.power(1 - np.max(cams, axis=0, keepdims=True), bg_factor) | |
| cams = np.concatenate((bg_score, cams), axis=0) | |
| prob = cams | |
| image = image.astype(np.uint8).transpose(1, 2, 0) | |
| prob = self.postprocessor(image, prob) | |
| label = np.argmax(prob, axis=0) | |
| label_tensor = torch.from_numpy(label).long() | |
| refined_mask = F.one_hot(label_tensor).to(device=self.device) | |
| refined_mask = refined_mask.permute(2, 0, 1) | |
| refined_mask = refined_mask[1:].float() | |
| return refined_mask | |
| def __call__(self, image, cams, separate=False, bg_factor=1.0): | |
| mean_bgr = (104.008, 116.669, 122.675) | |
| # covert Image to numpy array | |
| image = np.array(image).astype(np.float32) | |
| # RGB -> BGR | |
| image = image[:, :, ::-1] | |
| # Mean subtraction | |
| image -= mean_bgr | |
| # HWC -> CHW | |
| image = image.transpose(2, 0, 1) | |
| if isinstance(cams, torch.Tensor): | |
| cams = cams.cpu().detach().numpy() | |
| if separate: | |
| refined_mask = [ | |
| self.apply_crf(image, cam[None], bg_factor) for cam in cams | |
| ] | |
| refined_mask = torch.cat(refined_mask, dim=0) | |
| else: | |
| refined_mask = self.apply_crf(image, cams, bg_factor) | |
| return refined_mask | |