| from pathlib import Path | |
| import albumentations as A | |
| import cv2 | |
| import torch | |
| import tqdm | |
| from albumentations.pytorch.functional import img_to_tensor | |
| from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD | |
| from torchvision.utils import draw_segmentation_masks, make_grid, save_image | |
| import utils.misc as misc | |
| from models import get_ensemble_model | |
| from opt import get_opt | |
| def demo(folder_path, output_path=Path("tmp")): | |
| opt = get_opt() | |
| model = get_ensemble_model(opt).to(opt.device) | |
| misc.resume_from(model, opt.resume) | |
| with torch.no_grad(): | |
| for image_path in tqdm.tqdm(folder_path.glob("*.jpg")): | |
| image = cv2.imread(image_path.as_posix()) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| dsm_image = torch.from_numpy(image).permute(2, 0, 1) | |
| image_size = image.shape[:2] | |
| raw_image = img_to_tensor(image) | |
| image = img_to_tensor( | |
| image, | |
| normalize={"mean": IMAGENET_DEFAULT_MEAN, "std": IMAGENET_DEFAULT_STD}, | |
| ) | |
| image = image.to(opt.device).unsqueeze(0) | |
| outputs = model(image, seg_size=image_size) | |
| out_map = outputs["ensemble"]["out_map"][0, ...].detach().cpu() | |
| pred = outputs["ensemble"]["out_map"].max().item() | |
| if pred > opt.mask_threshold: | |
| print(f"Found manipulation in {image_path.name}") | |
| else: | |
| print(f"No manipulation found in {image_path.name}") | |
| overlay = draw_segmentation_masks( | |
| dsm_image, masks=out_map[0, ...] > opt.mask_threshold | |
| ) | |
| grid_image = make_grid( | |
| [ | |
| raw_image, | |
| (out_map.repeat(3, 1, 1) > opt.mask_threshold).float() * 255, | |
| overlay / 255.0, | |
| ], | |
| padding=5, | |
| ) | |
| image_name = image_path.stem + f"-{pred:.2f}" + image_path.suffix | |
| save_image(grid_image, (output_path / image_name).as_posix()) | |
| if __name__ == "__main__": | |
| folder_path = Path("demo") | |
| output_path = Path("tmp") | |
| output_path.mkdir(exist_ok=True, parents=True) | |
| demo(folder_path) | |