Skip to content

add stereotb option and support code #133

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 14, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions scripts/depthmap.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@ def ui(self, is_img2img):
with gr.Group():
with gr.Row():
gen_stereo = gr.Checkbox(label="Generate Stereo side-by-side image",value=False)
gen_stereotb = gr.Checkbox(label="Generate Stereo top-bottom image",value=False)
gen_anaglyph = gr.Checkbox(label="Generate Stereo anaglyph image (red/cyan)",value=False)
with gr.Row():
stereo_divergence = gr.Slider(minimum=0.05, maximum=10.005, step=0.01, label='Divergence (3D effect)', value=2.5)
Expand Down Expand Up @@ -137,10 +138,10 @@ def ui(self, is_img2img):
outputs=[clipthreshold_far]
)

return [compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, background_removal_model, background_removal, pre_depth_background_removal, save_background_removal_masks, gen_normal]
return [compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_stereotb, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, background_removal_model, background_removal, pre_depth_background_removal, save_background_removal_masks, gen_normal]

# run from script in txt2img or img2img
def run(self, p, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, background_removal_model, background_removal, pre_depth_background_removal, save_background_removal_masks, gen_normal):
def run(self, p, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_stereotb, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, background_removal_model, background_removal, pre_depth_background_removal, save_background_removal_masks, gen_normal):

# sd process
processed = processing.process_images(p)
Expand All @@ -163,14 +164,14 @@ def run(self, p, compute_device, model_type, net_width, net_height, match_size,
else:
background_removed_images = batched_background_removal(inputimages, background_removal_model)

newmaps, mesh_fi = run_depthmap(processed, p.outpath_samples, inputimages, None, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, "mp4", 0, background_removal, background_removed_images, save_background_removal_masks, gen_normal)
newmaps, mesh_fi = run_depthmap(processed, p.outpath_samples, inputimages, None, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_stereotb, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, "mp4", 0, background_removal, background_removed_images, save_background_removal_masks, gen_normal)

for img in newmaps:
processed.images.append(img)

return processed

def run_depthmap(processed, outpath, inputimages, inputnames, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, fnExt, vid_ssaa, background_removal, background_removed_images, save_background_removal_masks, gen_normal):
def run_depthmap(processed, outpath, inputimages, inputnames, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_stereotb, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, fnExt, vid_ssaa, background_removal, background_removed_images, save_background_removal_masks, gen_normal):

if len(inputimages) == 0 or inputimages[0] == None:
return []
Expand Down Expand Up @@ -461,7 +462,7 @@ def run_depthmap(processed, outpath, inputimages, inputnames, compute_device, mo
heatmap = (colormap(img_output2[:,:,0] / 256.0) * 2**16).astype(np.uint16)[:,:,:3]
outimages.append(heatmap)

if gen_stereo or gen_anaglyph:
if gen_stereo or gen_stereotb or gen_anaglyph:
print("Generating Stereo image..")
#img_output = cv2.blur(img_output, (3, 3))
balance = (stereo_balance + 1) / 2
Expand All @@ -471,27 +472,35 @@ def run_depthmap(processed, outpath, inputimages, inputnames, compute_device, mo
right_image = original_image if balance > 0.999 else \
apply_stereo_divergence(original_image, img_output, stereo_divergence * (1 - balance), stereo_fill)
stereo_img = np.hstack([left_image, right_image])
stereotb_img = np.vstack([left_image, right_image])

# flip sbs left/right if enabled in settings
if hasattr(opts, 'depthmap_script_sbsflip'):
if opts.depthmap_script_sbsflip:
stereo_img = np.hstack([right_image, left_image])
stereotb_img = np.vstack([right_image, left_image])

if gen_stereo:
outimages.append(stereo_img)
if gen_stereotb:
outimages.append(stereotb_img)
if gen_anaglyph:
print("Generating Anaglyph image..")
anaglyph_img = overlap(left_image, right_image)
outimages.append(anaglyph_img)
if (processed is not None):
if gen_stereo:
images.save_image(Image.fromarray(stereo_img), outpath, "", processed.all_seeds[count], processed.all_prompts[count], opts.samples_format, info=info, p=processed, suffix="_stereo")
if gen_stereotb:
images.save_image(Image.fromarray(stereotb_img), outpath, "", processed.all_seeds[count], processed.all_prompts[count], opts.samples_format, info=info, p=processed, suffix="_stereotb")
if gen_anaglyph:
images.save_image(Image.fromarray(anaglyph_img), outpath, "", processed.all_seeds[count], processed.all_prompts[count], opts.samples_format, info=info, p=processed, suffix="_anaglyph")
else:
# from tab
if gen_stereo:
images.save_image(Image.fromarray(stereo_img), 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, suffix="_stereo")
if gen_stereotb:
images.save_image(Image.fromarray(stereotb_img), 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, suffix="_stereotb")
if gen_anaglyph:
images.save_image(Image.fromarray(anaglyph_img), 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, suffix="_anaglyph")

Expand Down Expand Up @@ -1079,6 +1088,7 @@ def run_generate(depthmap_mode,
combine_output,
combine_output_axis,
gen_stereo,
gen_stereotb,
gen_anaglyph,
stereo_divergence,
stereo_fill,
Expand Down Expand Up @@ -1143,7 +1153,7 @@ def run_generate(depthmap_mode,
else:
background_removed_images = batched_background_removal(imageArr, background_removal_model)

outputs, mesh_fi = run_depthmap(None, outpath, imageArr, imageNameArr, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, fnExt, vid_ssaa, background_removal, background_removed_images, save_background_removal_masks, False)
outputs, mesh_fi = run_depthmap(None, outpath, imageArr, imageNameArr, compute_device, model_type, net_width, net_height, match_size, invert_depth, boost, save_depth, show_depth, show_heat, combine_output, combine_output_axis, gen_stereo, gen_stereotb, gen_anaglyph, stereo_divergence, stereo_fill, stereo_balance, clipdepth, clipthreshold_far, clipthreshold_near, inpaint, inpaint_vids, fnExt, vid_ssaa, background_removal, background_removed_images, save_background_removal_masks, False)

return outputs, mesh_fi, plaintext_to_html('info'), ''

Expand Down Expand Up @@ -1199,6 +1209,7 @@ def on_ui_tabs():
with gr.Group():
with gr.Row():
gen_stereo = gr.Checkbox(label="Generate Stereo side-by-side image",value=False)
gen_stereotb = gr.Checkbox(label="Generate Stereo top-bottom image",value=False)
gen_anaglyph = gr.Checkbox(label="Generate Stereo anaglyph image (red/cyan)",value=False)
with gr.Row():
stereo_divergence = gr.Slider(minimum=0.05, maximum=10.005, step=0.01, label='Divergence (3D effect)', value=2.5)
Expand Down Expand Up @@ -1284,6 +1295,7 @@ def on_ui_tabs():
combine_output,
combine_output_axis,
gen_stereo,
gen_stereotb,
gen_anaglyph,
stereo_divergence,
stereo_fill,
Expand Down