| 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 | |