| import os |
| import random |
| import argparse |
| import numpy as np |
| from PIL import Image |
| from pathlib import Path |
| from skimage import color |
| from tqdm.auto import tqdm |
|
|
| from lib.lime import LIME |
| from lib.fog_gen import fogAttenuation |
|
|
| from lib.gen_utils import ( |
| illumination2opacity, |
| reduce_lightHSV, |
| scale_depth) |
|
|
|
|
| def parse_arguments(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--clear_path", type=str, required=True, help="path to the file or the folder") |
| parser.add_argument("--depth_path", type=str, required=True, help="path to the file or the folder") |
| parser.add_argument("--save_folder", type=str, default="./generated/", help="path to the folder") |
| parser.add_argument("--txt_file", default=None, help="path to the folder") |
| parser.add_argument("--show", action="store_true") |
| return parser.parse_args() |
|
|
|
|
|
|
| class FogEffectGenerator: |
| def __init__(self): |
| self._lime = LIME(iterations=25, alpha=1.0) |
| |
| self._illumination2darkness = {0: 1, 1: 0.9, 2: 0.8, 3: 0.7} |
| self._weather2visibility = (500, 2000) |
| |
| self._illumination2fogcolor = {0: (150, 180), 1: (180, 200), 2: (200, 240), 3: (200, 240)} |
| |
| def getIlluminationMap(self, img: np.ndarray) -> np.ndarray: |
| self._lime.load(img) |
| T = self._lime.illumMap() |
| return T |
|
|
| def getIlluminationMapCheat(self, img: np.ndarray) -> np.ndarray: |
| T = color.rgb2gray(img) |
| return T |
| |
| def genEffect(self, img_path: str, depth_path: str): |
| I = np.array(Image.open(img_path)) |
| D = np.load(depth_path) |
|
|
| hI, wI, _ = I.shape |
| hD, wD = D.shape |
| |
| if hI!=hD or wI!=wD: |
| D = scale_depth(D, hI, wI) |
| |
| |
| T = self.getIlluminationMapCheat(I) |
| illumination_array = np.histogram(T, bins=4, range=(0,1))[0]/(T.size) |
| illumination = illumination_array.argmax() |
| |
| if illumination>0: |
| vmax = self._weather2visibility[1] if self._weather2visibility[1]<=D.max() else D.max() |
| if vmax<= self._weather2visibility[0]: |
| visibility = self._weather2visibility[0] |
| else: |
| visibility = random.randint(self._weather2visibility[0], int(vmax)) |
| fog_color = random.randint(self._illumination2fogcolor[illumination][0], self._illumination2fogcolor[illumination][1]) |
| I_dark = reduce_lightHSV(I, sat_red=self._illumination2darkness[illumination], val_red=self._illumination2darkness[illumination]) |
| I_fog = fogAttenuation(I_dark, D, visibility=visibility, fog_color=fog_color) |
| else: |
| fog_color = 75 |
| visibility = 150 |
| I_fog = fogAttenuation(I, D, visibility=visibility, fog_color=fog_color) |
| |
| return I_fog |
|
|
| def main(): |
| args = parse_arguments() |
| foggen = FogEffectGenerator() |
| |
| clearP = Path(args.clear_path) |
| depthP = Path(args.depth_path) |
| if clearP.is_file() and (depthP.is_file() and depthP.suffix==".npy"): |
| snowy = foggen.genEffect(clearP, depthP) |
| if args.show: |
| Image.fromarray(snowy).show() |
|
|
| if clearP.is_dir() and depthP.is_dir(): |
| if args.txt_file: |
| with open(args.txt_file, 'r') as f: |
| files = f.read().split('\n') |
| image_files = [clearP / f for f in files] |
| else: |
| image_files = sorted(Path(clearP).glob("*")) |
|
|
| depth_files = [Path(depthP) / ("-".join(imgf.name.split('-')[:2])+".npy") for imgf in image_files] |
| |
| valid_files = [idx for idx, f in enumerate(depth_files) if f.exists()] |
| image_files = [image_files[idx] for idx in valid_files] |
| depth_files = [depth_files[idx] for idx in valid_files] |
| |
| save_folder = Path(args.save_folder) |
| if not save_folder.exists(): |
| os.makedirs(str(save_folder)) |
| |
| for imgp, depthp in tqdm(zip(image_files, depth_files), total=len(image_files)): |
| foggy = foggen.genEffect(imgp, depthp) |
| Image.fromarray(foggy).save(save_folder / (imgp.stem+"-fsyn.jpg")) |
|
|
| if __name__=='__main__': |
| main() |
