Introduce separation parameter for stereo image generation

Author: semjonsona

scripts/depthmap.py

@@ -147,6 +147,8 @@ def main_ui_panel(is_depth_tab):
 			with gr.Row(visible=False) as stereo_options_row_1:
 				stereo_divergence = gr.Slider(minimum=0.05, maximum=10.005, step=0.01, label='Divergence (3D effect)',
 											  value=2.5)
+				stereo_separation = gr.Slider(minimum=-5.0, maximum=5.0, step=0.01, label='Separation (moves images apart)',
+											  value=0.0)
 			with gr.Row(visible=False) as stereo_options_row_2:
 				stereo_fill = gr.Dropdown(label="Gap fill technique",
 										  choices=['none', 'naive', 'naive_interpolating', 'polylines_soft',
@@ -266,7 +268,7 @@ def background_removal_options_visibility(v):
 			outputs=[bgrem_options_row_1, bgrem_options_row_2]
 		)
 
-	return [compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical]
+	return [compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical]
 
 
 class Script(scripts.Script):
@@ -283,7 +285,7 @@ def ui(self, is_img2img):
 
 	# run from script in txt2img or img2img
 	def run(self, p,
-			compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical
+			compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical
 			):
 
 		# sd process 
@@ -309,7 +311,7 @@ def run(self, p,
 
 		newmaps, mesh_fi, meshsimple_fi = run_depthmap(processed, p.outpath_samples, inputimages, None,
                                         compute_device, model_type,
-                                        net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
+                                        net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
                                         background_removed_images, "mp4", 0, False, None, False, gen_mesh, mesh_occlude, mesh_spherical )
 
 		for img in newmaps:
@@ -328,7 +330,7 @@ def reload_sd_model():
 		shared.sd_model.first_stage_model.to(devices.device)
 
 def run_depthmap(processed, outpath, inputimages, inputnames,
-                 compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
+                 compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
                  background_removed_images, fnExt, vid_ssaa, custom_depthmap, custom_depthmap_img, depthmap_batch_reuse, gen_mesh, mesh_occlude, mesh_spherical):
 
 	if len(inputimages) == 0 or inputimages[0] == None:
@@ -687,7 +689,7 @@ def run_depthmap(processed, outpath, inputimages, inputnames,
 							try:
 								images.save_image(Image.fromarray(img_output), outpath, "", processed.all_seeds[count], processed.all_prompts[count], opts.samples_format, info=info, p=processed, suffix="_depth")
 							except ValueError as ve:
-								if not 'image has wrong mode' in str(ve): raise ve
+								if not ('image has wrong mode' in str(ve) or 'cannot write mode I;16 as JPEG' in str(ve)): raise ve
 						else:
 							images.save_image(Image.fromarray(img_output2), outpath, "", processed.all_seeds[count], processed.all_prompts[count], opts.samples_format, info=info, p=processed, suffix="_depth")
 					elif save_depth:
@@ -697,7 +699,7 @@ def run_depthmap(processed, outpath, inputimages, inputnames,
 							try:
 								images.save_image(Image.fromarray(img_output), path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True,no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=None, forced_filename=None)
 							except ValueError as ve:
-								if not 'image has wrong mode' in str(ve): raise ve
+								if not ('image has wrong mode' in str(ve) or 'cannot write mode I;16 as JPEG' in str(ve)): raise ve
 						else:
 							images.save_image(Image.fromarray(img_output2), path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True,no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=None, forced_filename=None)
 				else:
@@ -714,7 +716,7 @@ def run_depthmap(processed, outpath, inputimages, inputnames,
 				print("Generating stereoscopic images..")
 
 				stereomodes = stereo_modes
-				stereoimages = create_stereoimages(inputimages[count], img_output, stereo_divergence, stereomodes, stereo_balance, stereo_fill)
+				stereoimages = create_stereoimages(inputimages[count], img_output, stereo_divergence, stereo_separation, stereomodes, stereo_balance, stereo_fill)
 
 				for c in range(0, len(stereoimages)):
 					outimages.append(stereoimages[c])
@@ -1138,6 +1140,7 @@ def run_generate(depthmap_mode,
                 gen_stereo,
                 stereo_modes,
                 stereo_divergence,
+				stereo_separation,
                 stereo_fill,
                 stereo_balance,
                 inpaint,
@@ -1205,7 +1208,7 @@ def run_generate(depthmap_mode,
 
 	outputs, mesh_fi, meshsimple_fi = run_depthmap(
         None, outpath, imageArr, imageNameArr,
-        compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
+        compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal,
         background_removed_images, fnExt, vid_ssaa, custom_depthmap, custom_depthmap_img, depthmap_batch_reuse, gen_mesh, mesh_occlude, mesh_spherical)
 
 	# use inpainted 3d mesh to show in 3d model output when enabled in settings
@@ -1264,7 +1267,7 @@ def on_ui_tabs():
                 submit = gr.Button('Generate', elem_id="depthmap_generate", variant='primary')
 
 				# insert main panel
-                compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical = main_ui_panel(True)
+                compute_device, model_type, net_width, net_height, match_size, boost, invert_depth, clipdepth, clipthreshold_far, clipthreshold_near, combine_output, combine_output_axis, save_depth, show_depth, show_heat, gen_stereo, stereo_modes, stereo_divergence, stereo_separation, stereo_fill, stereo_balance, inpaint, inpaint_vids, background_removal, save_background_removal_masks, gen_normal, pre_depth_background_removal, background_removal_model, gen_mesh, mesh_occlude, mesh_spherical = main_ui_panel(True)
 
                 unloadmodels = gr.Button('Unload models', elem_id="depthmap_unloadmodels")
 
@@ -1357,6 +1360,7 @@ def custom_depthmap_visibility(v):
 				gen_stereo,
 				stereo_modes,
 				stereo_divergence,
+				stereo_separation,
 				stereo_fill,
 				stereo_balance,
 				inpaint,

scripts/stereoimage_generation.py

@@ -2,7 +2,7 @@
 import numpy as np
 from PIL import Image
 
-def create_stereoimages(original_image, depthmap, divergence, modes=None, stereo_balance=0.0,
+def create_stereoimages(original_image, depthmap, divergence, separation=0.0, modes=None, stereo_balance=0.0,
                         fill_technique='polylines_sharp'):
     """Creates stereoscopic images.
     An effort is made to make them look nice, but beware that the resulting image will have some distortion .
@@ -11,6 +11,9 @@ def create_stereoimages(original_image, depthmap, divergence, modes=None, stereo
     :param depthmap: depthmap corresponding to the original image. White = near, black = far.
     :param float divergence: the measure of 3D effect, in percentages.
       A good value will likely be somewhere in the [0.05; 10.0) interval.
+    :param float separation: measure by how much to move two halfs of the spereoimage apart from eachother.
+      Measured in percentages. Negative values move two parts closer togethert.
+      Affects which parts of the image will be visible in left and/or right half.
     :param list modes: how the result will look like. By default only 'left-right' is generated
       - a picture for the left eye will be on the left and the picture from the right eye - on the rigth.
       The supported modes are: 'left-right', 'right-left', 'top-bottom', 'bottom-top', 'red-cyan-anaglyph'.
@@ -28,9 +31,9 @@ def create_stereoimages(original_image, depthmap, divergence, modes=None, stereo
     original_image = np.asarray(original_image)
     balance = (stereo_balance + 1) / 2
     left_eye = original_image if balance < 0.001 else \
-        apply_stereo_divergence(original_image, depthmap, +1 * divergence * balance, fill_technique)
+        apply_stereo_divergence(original_image, depthmap, +1 * divergence * balance, -1 * separation, fill_technique)
     right_eye = original_image if balance > 0.999 else \
-        apply_stereo_divergence(original_image, depthmap, -1 * divergence * (1 - balance), fill_technique)
+        apply_stereo_divergence(original_image, depthmap, -1 * divergence * (1 - balance), separation, fill_technique)
 
     results = []
     for mode in modes:
@@ -49,20 +52,26 @@ def create_stereoimages(original_image, depthmap, divergence, modes=None, stereo
     return [Image.fromarray(r) for r in results]
 
 
-def apply_stereo_divergence(original_image, depth, divergence, fill_technique):
+def apply_stereo_divergence(original_image, depth, divergence, separation, fill_technique):
     depth_min = depth.min()
     depth_max = depth.max()
     normalized_depth = (depth - depth_min) / (depth_max - depth_min)
     divergence_px = (divergence / 100.0) * original_image.shape[1]
+    separation_px = (separation / 100.0) * original_image.shape[1]
 
     if fill_technique in ['none', 'naive', 'naive_interpolating']:
-        return apply_stereo_divergence_naive(original_image, normalized_depth, divergence_px, fill_technique)
+        return apply_stereo_divergence_naive(
+            original_image, normalized_depth, divergence_px, separation_px, fill_technique
+        )
     if fill_technique in ['polylines_soft', 'polylines_sharp']:
-        return apply_stereo_divergence_polylines(original_image, normalized_depth, divergence_px, fill_technique)
+        return apply_stereo_divergence_polylines(
+            original_image, normalized_depth, divergence_px, separation_px, fill_technique
+        )
 
 
 @njit
-def apply_stereo_divergence_naive(original_image, normalized_depth, divergence_px: float, fill_technique):
+def apply_stereo_divergence_naive(
+        original_image, normalized_depth, divergence_px: float, separation_px: float, fill_technique):
     h, w, c = original_image.shape
 
     derived_image = np.zeros_like(original_image)
@@ -72,7 +81,7 @@ def apply_stereo_divergence_naive(original_image, normalized_depth, divergence_p
         # Swipe order should ensure that pixels that are closer overwrite
         # (at their destination) pixels that are less close
         for col in range(w) if divergence_px < 0 else range(w - 1, -1, -1):
-            col_d = col + int((normalized_depth[row][col] ** 2) * divergence_px)
+            col_d = col + int((normalized_depth[row][col] ** 2) * divergence_px + separation_px)
             if 0 <= col_d < w:
                 derived_image[row][col_d] = original_image[row][col]
                 filled[row * w + col_d] = 1
@@ -127,7 +136,8 @@ def apply_stereo_divergence_naive(original_image, normalized_depth, divergence_p
 
 
 @njit(parallel=True)  # fastmath=True does not reasonably improve performance
-def apply_stereo_divergence_polylines(original_image, normalized_depth, divergence_px: float, fill_technique):
+def apply_stereo_divergence_polylines(
+        original_image, normalized_depth, divergence_px: float, separation_px: float, fill_technique):
     # This code treats rows of the image as polylines
     # It generates polylines, morphs them (applies divergence) to them, and then rasterizes them
     EPSILON = 1e-7
@@ -141,19 +151,19 @@ def apply_stereo_divergence_polylines(original_image, normalized_depth, divergen
         # format: new coordinate of the vertex, divergence (closeness), column of pixel that contains the point's color
         pt = np.zeros((5 + 2 * w, 3), dtype=np.float_)
         pt_end: int = 0
-        pt[pt_end] = [-3.0 * abs(divergence_px), 0.0, 0.0]
+        pt[pt_end] = [-1.0 * w, 0.0, 0.0]
         pt_end += 1
         for col in range(0, w):
             coord_d = (normalized_depth[row][col] ** 2) * divergence_px
-            coord_x = col + 0.5 + coord_d
+            coord_x = col + 0.5 + coord_d + separation_px
             if PIXEL_HALF_WIDTH < EPSILON:
                 pt[pt_end] = [coord_x, abs(coord_d), col]
                 pt_end += 1
             else:
                 pt[pt_end] = [coord_x - PIXEL_HALF_WIDTH, abs(coord_d), col]
                 pt[pt_end + 1] = [coord_x + PIXEL_HALF_WIDTH, abs(coord_d), col]
                 pt_end += 2
-        pt[pt_end] = [w + 3.0 * abs(divergence_px), 0.0, w - 1]
+        pt[pt_end] = [2.0 * w, 0.0, w - 1]
         pt_end += 1
 
         # generating the segments of the morphed polyline