diff --git a/.gitignore b/.gitignore index 9a62697..efa1300 100644 --- a/.gitignore +++ b/.gitignore @@ -180,4 +180,11 @@ checkpoints-legacy loss_single.sh normal_vs_hq.sh dataset-hq -generations \ No newline at end of file +generations +generate.sh +image_grid.py +deliberate.safetensors +generations-paper.tar.gz +test.sh +header_masks +header_masks.tar.gz \ No newline at end of file diff --git a/README.md b/README.md index 5782d52..c0112ae 100644 --- a/README.md +++ b/README.md @@ -102,3 +102,95 @@ python evaluate_multi_class.py \ } ] ``` + +### Commands used to generate the paper images + +```sh +python generate_and_segment.py \ + --prompt "a photograph of a bear in a pine forest observing an abandoned refrigerator" \ + --classes "bear,refrigerator" \ + --mask-rcnn-classes "bear,refrigerator" \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a photograph of a tennisplayer hitting a tennis ball with his racket" \ + --classes "racket,person" \ + --mask-rcnn-classes "person,tennis racket" \ + --n-images 5 +``` + +```sh +python generate_and_segment.py \ + --prompt "a turtle climbing on top of a reef in the pacific ocean. photorealistic digital art, cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, 3 d ue 5, 4 k, hq" \ + --classes "turtle" \ + --skip-mask-rcnn \ + --n-images 3 + +python generate_and_segment.py \ + --prompt "a flamingo standing next to a yacht. photorealistic digital art, cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, 3 d ue 5, 4 k, hq" \ + --classes "flamingo,yacht" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a long shot photograph of an elephant standing in front of the pyramids. photorealistic digital art, cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, 3 d ue 5, 4 k, hq" \ + --classes "elephant" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a long shot photograph of an elephant standing in front of an eurofighter jet. photorealistic digital art, cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, 3 d ue 5, 4 k, hq" \ + --classes "elephant,eurofighter" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a wooden barrel next to a cactus outside of a saloon in late 1800s austin texas. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "barrel,cactus" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a photograph of a plate of carbonara pasta. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "carbonara" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a photograph of a lamborghini murcielago. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "lamborghini" \ + --skip-mask-rcnn \ + --n-images 5 + +# Generated with the fine-tuned checkpoint +# https://civitai.com/models/4823/deliberate +python generate_and_segment.py \ + --prompt "a long shot photograph of a man standing in front of of a ferrari testarossa. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "ferrari,person" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a long shot photograph of lion chasing a zebra in the african savannah. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "lion,zebra" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a photograph of a sushi platter next to an ashai beer can on top of a table. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "sushi,beercan" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a photograph of an alpaca standing next to a park bench. cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "alpaca,parkbench" \ + --skip-mask-rcnn \ + --n-images 5 + +python generate_and_segment.py \ + --prompt "a portrait of an explorer dressed in the style of indiana jones wearing a hat in front of an egyptian sphinx. photorealistic digital art, cinematic, extremely high detail, cinematic lighting, trending, artstation, cgsociety, colorful, 4k, hq" \ + --classes "explorer" \ + --skip-mask-rcnn \ + --n-images 5 +``` diff --git a/generate_and_segment.py b/generate_and_segment.py index f25facd..a0dd344 100644 --- a/generate_and_segment.py +++ b/generate_and_segment.py @@ -20,11 +20,15 @@ parser = argparse.ArgumentParser(prog="grounded generation") -parser.add_argument("--use-sd2", action="store_true") +parser.add_argument("--use-sd2", action=argparse.BooleanOptionalAction) +parser.add_argument("--skip-mask-rcnn", action=argparse.BooleanOptionalAction) parser.add_argument("--output-dir", type=str, default="generations") parser.add_argument("--prompt", type=str, default="a photograph of a cat and a dog") +parser.add_argument("--negative-prompt", type=str, default=None) parser.add_argument("--classes", type=str, default="cat,dog") -parser.add_argument("--grounding-ckpt", type=str, default="checkpoints/normal_arch_checkpoint.pth") +parser.add_argument("--mask-rcnn-classes", type=str, default="cat,dog") +parser.add_argument("--grounding-ckpt", type=str, default="mmdetection/checkpoint/grounding_module.pth") +parser.add_argument("--n-images", type=int, default=1) parser.add_argument("--seed", type=int, default=2147483647) args = parser.parse_args() @@ -44,7 +48,9 @@ # Load COCO classes coco_classes = open("mmdetection/demo/coco_80_class.txt").read().split("\n") -coco_classes = dict([(i, c) for i, c in enumerate(coco_classes)]) + +class2label = dict([(i, c) for i, c in enumerate(coco_classes)]) +label2class = dict([(c, i) for i, c in enumerate(coco_classes)]) # Add the current timestamp to the output folder args.output_dir = os.path.join(args.output_dir, str(int(time.time()))) @@ -86,94 +92,120 @@ with torch.no_grad(): classes = [c.strip() for c in args.classes.split(",")] - print(f"Generating sample using prompt: {args.prompt}") - print(f"The target classes are: {classes}") - - grounded_unet.clear_grounding_features() - - # Sample an image - image = pipeline(args.prompt, generator=rand_generator).images[0] - array_image = np.array(image) - - # Get the Mask R-CNN segmentation - _, mask_rcnn_segmentations = inference_detector( - pretrain_detector, - [array_image] - ).pop() - - # Save all masks from Mask R-CNN - for i, masks in enumerate(mask_rcnn_segmentations): - if len(masks) == 0: - continue - - masked_image = Image.fromarray(plot_mask( - np.array(image), - np.expand_dims(masks[0], 0) - )) - - masked_image.save( - os.path.join( - args.output_dir, - f"masked_image_{coco_classes[i]}_mask_rcnn.png" + mask_rcnn_classes = [c.strip() for c in args.mask_rcnn_classes.split(",")] + mask_rcnn_classes = [c for c in mask_rcnn_classes if c in label2class] + + for i in range(args.n_images): + image_dir = os.path.join(args.output_dir, str(i)) + os.makedirs(image_dir, exist_ok=True) + + print(f"Generating sample {i} using prompt: {args.prompt}") + print(f"The target classes are: {classes}") + + grounded_unet.clear_grounding_features() + + # Sample an image + image = pipeline( + prompt=args.prompt, + negative_prompt=args.negative_prompt, + generator=rand_generator + ).images[0] + + array_image = np.array(image) + + # Get the Mask R-CNN segmentation + _, mask_rcnn_segmentations = inference_detector( + pretrain_detector, + [array_image] + ).pop() + + if not args.skip_mask_rcnn: + picked_rcnn_masks = [] + + # Save all masks from Mask R-CNN + for i, masks in enumerate(mask_rcnn_segmentations): + if len(masks) == 0: + continue + + masked_image = Image.fromarray(plot_mask( + np.array(image), + np.expand_dims(masks[0], 0) + )) + + mask_rcnn_label = class2label[i] + + if mask_rcnn_label in mask_rcnn_classes: + picked_rcnn_masks.append(masks[0]) + + masked_image.save( + os.path.join( + image_dir, + f"masked_image_{mask_rcnn_label}_mask_rcnn.png" + ) + ) + + # Visualize all the target Mask R-CNN masks at one + if len(picked_rcnn_masks) > 0: + all_mask_rcnn_image = Image.fromarray(plot_mask(np.array(image), picked_rcnn_masks)) + all_mask_rcnn_image.save(os.path.join(image_dir, f"mask_rcnn_all_masks.png")) + + # Save the genearted image + image.save(os.path.join(image_dir, f"sd_image.png")) + + # Get the UNet features + unet_features = grounded_unet.get_grounding_features() + + # Extract embeddings for each individual class, + # using the prompt "a photograph of a {x}" + single_class_embeddings = {} + + for label in classes: + single_class_embeddings[label] = get_embeddings( + tokenizer=tokenizer, + embedder=embedder, + device=device, + prompt=label, + # prompt=f"a photograph of a {label}", + labels=[label], + inverted_vocab=tokenizer_inverted_vocab + )[label] + + all_fusion_masks = [] + + for label in classes: + embedding = single_class_embeddings[label] + + # Subtract the embeddings from all other classes + for other_label in set(classes) - set([label]): + embedding -= single_class_embeddings[other_label] + + fusion_segmentation = seg_module(unet_features, embedding) + fusion_segmentation_pred = fusion_segmentation[0, 0, :, :] + fusion_mask = preprocess_mask(mask=fusion_segmentation_pred.unsqueeze(0)) + + # Save the fusion mask + torchvision.utils.save_image( + torch.from_numpy(fusion_mask), + os.path.join(image_dir, f"mask_{label}_segmodule.png"), + normalize=True, + scale_each=True, ) - ) - - # Save the genearted image - image.save(os.path.join(args.output_dir, f"sd_image.png")) - - # Get the UNet features - unet_features = grounded_unet.get_grounding_features() - # Extract embeddings for each individual class, - # using the prompt "a photograph of a {x}" - single_class_embeddings = {} + # Also plot the mask over the image + masked_image = Image.fromarray(plot_mask(np.array(image), np.expand_dims(fusion_mask, 0))) + masked_image.save(os.path.join(image_dir, f"masked_image_{label}_segmodule.png")) - for label in classes: - single_class_embeddings[label] = get_embeddings( - tokenizer=tokenizer, - embedder=embedder, - device=device, - prompt=f"a photograph of a {label}", - labels=[label], - inverted_vocab=tokenizer_inverted_vocab - )[label] + all_fusion_masks.append(fusion_mask) - all_fusion_masks = [] + # Mask the original image and save the cutted out portion + expanded_mask = np.stack([fusion_mask.astype(int)] * 3, axis=-1) - for label in classes: - embedding = single_class_embeddings[label] + extracted_image = np.array(image) + extracted_image[expanded_mask == 0] = 0 - # Subtract the embeddings from all other classes - for other_label in set(classes) - set([label]): - embedding -= single_class_embeddings[other_label] - - fusion_segmentation = seg_module(unet_features, embedding) - fusion_segmentation_pred = fusion_segmentation[0, 0, :, :] - fusion_mask = preprocess_mask(mask=fusion_segmentation_pred.unsqueeze(0)) - - # Save the fusion mask - torchvision.utils.save_image( - torch.from_numpy(fusion_mask), - os.path.join(args.output_dir, f"mask_{label}_segmodule.png"), - normalize=True, - scale_each=True, - ) - - # Also plot the mask over the image - masked_image = Image.fromarray(plot_mask(np.array(image), np.expand_dims(fusion_mask, 0))) - masked_image.save(os.path.join(args.output_dir, f"masked_image_{label}_segmodule.png")) - - all_fusion_masks.append(fusion_mask) - - # Mask the original image and save the cutted out portion - expanded_mask = np.stack([fusion_mask.astype(int)] * 3, axis=-1) - - extracted_image = np.array(image) - extracted_image[expanded_mask == 0] = 0 - - Image.fromarray(extracted_image).save( - os.path.join(args.output_dir, f"extracted_{label}_segmodule.png") - ) + Image.fromarray(extracted_image).save( + os.path.join(image_dir, f"extracted_{label}_segmodule.png") + ) - all_fusion_image = Image.fromarray(plot_mask(np.array(image), all_fusion_masks)) - all_fusion_image.save(os.path.join(args.output_dir, f"segmodule_all_masks.png")) + all_fusion_image = Image.fromarray(plot_mask(np.array(image), all_fusion_masks)) + all_fusion_image.save(os.path.join(image_dir, f"segmodule_all_masks.png"))