Skip to content

Commit

Permalink
Add SSIM blur slider to avoid detail clutter bogging down ssmi
Browse files Browse the repository at this point in the history
prompt_travel.py

- Add a new SSIM blur slider to the script's prompt menu.
- The slider helps with images featuring many small changing details.
- The value of the slider ranges from 0 to 100 with a step of 1.
- Update the script's configuration list and arguments with the new
slider.
- Apply the GaussianBlur filter to the images based on the slider value.
- Use the blurred images for calculating distance in the script's main
loop.
- Enhance the script's functionality by considering the SSIM blur when
checking distances between images.
  • Loading branch information
Daniel Dengler committed Nov 14, 2023
1 parent 2493c43 commit 66f7c4c
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions scripts/prompt_travel.py
Original file line number Diff line number Diff line change
Expand Up @@ -485,6 +485,9 @@ def switch_genesis(genesis:str):
ssim_diff_min = gr.Slider(
label="SSIM min threshold", value=75, minimum=0, maximum=100, step=1
)
ssim_blur = gr.Slider(
label="SSIM blur (helps with images featuring many small changing details)", value=0, minimum=0, maximum=100, step=1
)

return [
mode, lerp_meth, replace_dim, replace_order,
Expand All @@ -493,7 +496,7 @@ def switch_genesis(genesis:str):
upscale_meth, upscale_ratio, upscale_width, upscale_height,
video_fmt, video_fps, video_pad, video_pick,
ext_video, ext_upscale, ext_depth, ssim_diff, ssim_ccrop,
substep_min, ssim_diff_min
substep_min, ssim_diff_min, ssim_blur
]

def run(self, p:Processing,
Expand All @@ -504,7 +507,7 @@ def run(self, p:Processing,
video_fmt:str, video_fps:float, video_pad:int, video_pick:str,
ext_video:bool, ext_upscale:bool, ext_depth:bool,
ssim_diff: float, ssim_ccrop:int,
substep_min:float, ssim_diff_min:int
substep_min:float, ssim_diff_min:int, ssim_blur:int
):

# enum lookup
Expand Down Expand Up @@ -593,6 +596,7 @@ def run(self, p:Processing,
self.ssim_ccrop = ssim_ccrop
self.substep_min = substep_min
self.ssim_diff_min = ssim_diff_min
self.ssim_blur = ssim_blur

def upscale_image_callback(params:ImageSaveParams):
params.image = upscale_image(params.image, p.width, p.height, upscale_meth, upscale_ratio, upscale_width, upscale_height)
Expand Down Expand Up @@ -901,9 +905,16 @@ def ssim_loop(
break
check = False
for i in range(done, len(prompt_images) - 1):
a_img: PILImage = prompt_images[i]
b_img: PILImage = prompt_images[i + 1]
if self.ssim_blur > 0:
from PIL import ImageFilter
a_img: PILImage = prompt_images[i].filter(ImageFilter.GaussianBlur(radius=self.ssim_blur))
b_img: PILImage = prompt_images[i + 1].filter(ImageFilter.GaussianBlur(radius=self.ssim_blur))

# Check distance between i and i+1
a = transform(prompt_images[i]).unsqueeze(0)
b = transform(prompt_images[i + 1]).unsqueeze(0)
a = transform(a_img).unsqueeze(0)
b = transform(b_img).unsqueeze(0)
d = ssim(a, b)

if d < ssim_diff and (dists[i + 1] - dists[i]) > substep_min:
Expand Down

0 comments on commit 66f7c4c

Please sign in to comment.