| from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation |
| import numpy as np |
| import cv2 |
| import os |
| import torch |
| from skimage.metrics import structural_similarity as ssim |
|
|
| class SKY: |
| def __init__(self): |
| self.device = ("cuda" if torch.cuda.is_available() else "cpu") |
| self.processor = AutoImageProcessor.from_pretrained("models/mask2former-swin-large-ade-semantic") |
| self.model = Mask2FormerForUniversalSegmentation.from_pretrained("models/mask2former-swin-large-ade-semantic").to(self.device) |
|
|
| def _SKY_mask(self,image_rgb): |
|
|
| inputs = self.processor(images=image_rgb, return_tensors="pt").to(self.device) |
|
|
| outputs = self.model(**inputs) |
| inputs.to("cpu") |
| del inputs |
| predicted_semantic_map = self.processor.post_process_semantic_segmentation(outputs, target_sizes=[image_rgb.shape[:2]])[0] |
| mask = predicted_semantic_map.cpu().numpy() |
| predicted_semantic_map.to("cpu") |
| del predicted_semantic_map |
|
|
| mask_np = (mask == 2) |
|
|
| if np.count_nonzero(mask_np)==0: |
| return None |
|
|
| return mask_np |
|
|
| def SKY_remove(self,image_rgb,gamma=None): |
|
|
|
|
| mask = self._SKY_mask(image_rgb) |
| if mask is None: |
| return image_rgb |
| binary_mask = np.uint8(mask) * 255 |
|
|
|
|
| if gamma: |
| binary_mask = cv2.GaussianBlur(~binary_mask, (15, 15), gamma) |
| binary_mask = ~binary_mask |
|
|
| image_bgra = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2RGBA) |
| image_bgra[:, :, 3] = ~binary_mask |
|
|
| return image_bgra |
|
|
|
|
| def __del__(self): |
| self.model = None |
| self.processor = None |
| del self.model |
| del self.processor |
| torch.cuda.empty_cache() |
| import gc |
| gc.collect() |
|
|
| if __name__ == "__main__": |
| pass |
|
|
