SYSUykLin
commented
2 weeks ago Hello: I assume the depth and segment map is [H, W]，because I saw the default parameters condition_type=one_hot. The depth map is the normalized distance from 3D point cloud to camera. shape: [H, W] The seg map is the labels of the 3D point cloud in NYU40 class. shape: [H, W] The parameters of this code I copy from scripts/generate_outputs.py depth.npy is the distance from 3D point cloud to camera. shape: [H, W] segmentation.jpg is the rgb semantic map. shape:[H, W]. Shown as follow:
And I write the code:depth_npy_path = "scripts/depth.npy"
segment_path = "scripts/segmentation.jpg"
controlnet_conditioning_scale = [3.5, 1.5]
base_model_path = "ckpts/stable-diffusion-2-1"
checkpoint_path = "controlnet_ckpts/layout_diffusion_hypersim_prompt_one_hot_multi_control_bs32_epoch24"
checkpoint_subfolder = "checkpoint-10900"
condition_type = "one_hot"
def generate_output(condition_images, controlnet, pipe, generator):
prompts = ["A good-looking kitchen"]
extra_kwargs = dict()
if pipe.unet.config.in_channels == 5:
if "depths" not in extra_kwargs:
scale_h = SCALE_H[args.dataset]
depths = [transforms.Resize(scale_h, interpolation=transforms.InterpolationMode.NEAREST_EXACT)(depth) for depth in depth_images]
extra_kwargs.update(depths=depths)
depth_inv_norm_smalls = [depth_inv_norm_transforms(transforms.ToTensor()(depth)[0])
for depth in extra_kwargs["depths"]]
height, width = extra_kwargs["depths"][0].shape
extra_kwargs.update(depth_inv_norm_smalls=depth_inv_norm_smalls, height=height, width=width)
if isinstance(pipe, StableDiffusionControlNetPipeline):
extra_kwargs.update(image=control_images, condition_type=condition_type,
indice_to_embedding=None,
controlnet_conditioning_scale=controlnet_conditioning_scale,
no_depth_cond=False)
guide_images: List[Image.Image] = pipe(prompts,
num_inference_steps=20,
latents=None,
generator=generator,
guidance_scale=7.5,
**extra_kwargs).images
"""
假设：depth map的格式应该是单纯的深度
segment map的格式应是one hot
"""
def prepare_images(seg_path, depth_path):
depth_map = np.load(depth_path)
nan_mask = np.isnan(depth_map)
depth_map[nan_mask] = np.nanmax(depth_map)
max_depth = np.nanmax(depth_map)
min_depth = np.nanmin(depth_map)
norm = colors.LogNorm(vmin=min_depth, vmax=max_depth)
depth_map_norm = norm(depth_map)
# depth_map_norm = np.concatenate([np.array(depth_map_norm)[:, :, None]] * 3, axis=2).astype(np.uint8)
depth_images = [Image.fromarray(depth_map_norm)]
seg_map = cv2.imread(seg_path)
seg_map = cv2.cvtColor(seg_map, cv2.COLOR_BGR2RGB)
metadata_semantic_colors_hdf5_file = "scripts/metadata_semantic_colors.hdf5"
with h5py.File(metadata_semantic_colors_hdf5_file, "r") as f:
semantic_colors = f["dataset"][:]
image_idx = np.zeros((seg_map.shape[0], seg_map.shape[1], 1), dtype=np.uint8)
for i, color in enumerate(semantic_colors):
diff = np.mean(np.abs(seg_map[:, :, :] - color), axis=-1)
match_mask = (diff <= 10)
image_idx[match_mask] = i
image_idx = np.squeeze(image_idx)
seg_images = [Image.fromarray(image_idx.astype(np.uint8))]
# 都是[H, W]
return [seg_images, depth_images]
def load_ckpts():
unet: UNet2DConditionModel = UNet2DConditionModel.from_pretrained(
base_model_path, subfolder="unet", torch_dtype=torch.float16
)
if checkpoint_path is None:
controlnet = None
elif "multi_control" not in checkpoint_path:
print("single controlnet...")
controlnet = ControlNetModel.from_pretrained(
checkpoint_path, subfolder=f"{checkpoint_subfolder.strip('/')}/controlnetmodel", torch_dtype=torch.float16)
else:
print("multiple controlnet...")
controlnet: MultiControlNetModel = MultiControlNetModel.from_pretrained(
checkpoint_path, subfolder=f"{checkpoint_subfolder.strip('/')}", torch_dtype=torch.float16
)
assert isinstance(controlnet_conditioning_scale, (list, tuple))
if controlnet is not None:
pipe_method = StableDiffusionControlNetPipeline
pipe = pipe_method.from_pretrained(
base_model_path, unet=unet, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None,
)
else:
pipe = StableDiffusionPipeline.from_pretrained(
base_model_path, unet=unet, torch_dtype=torch.float16, safety_checker=None,
)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to(torch_device=torch.device('cuda'))
# pipe.enable_xformers_memory_efficient_attention()
pipe.set_progress_bar_config(disable=True)
generator = torch.Generator(device=pipe.device).manual_seed(10000)
return controlnet, pipe, generator
if __name__ == "__main__":
control_images = prepare_images(segment_path, depth_npy_path)
controlnet, pipe, generator = load_ckpts()
generate_output(control_images, controlnet, pipe, generator)the shape of depth map and seg map is [H, W]. But it show an error:
raise ValueError(f"Wrong channels of input condition, expected {self.config.conditioning_channels}, "
ValueError: Wrong channels of input condition, expected 8, got 3.Could you please tell me the format for inputting depth and semantic maps? Thank you
OrangeSodahub
Hello: great work!! I want to run on my own dataset. But I don't know the type of the segmentation map and depth map. For example, I have seg map and depth map:
What kind of transformation should I perform to make the depth map and segmentation map correctly input into your model and run properly? For example, the type of the depth map is gray map, or seg map is one hot or rgb type the same as the hypersim. Thanks