SAMControlNet/SyntheticDataSAM
2
1import gradio as gr2import numpy as np3import torch4import jax5import jax.numpy as jnp6from flax.jax_utils import replicate7from flax.training.common_utils import shard8from PIL import Image9from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator10from diffusers import (11 FlaxStableDiffusionControlNetPipeline,12 FlaxControlNetModel,13)14from transformers import pipeline15 16import colorsys17 18sam_checkpoint = "sam_vit_h_4b8939.pth"19model_type = "vit_h"20device = "cuda" if torch.cuda.is_available() else "cpu"21 22 23#sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)24#sam.to(device=device)25#predictor = SamPredictor(sam)26#mask_generator = SamAutomaticMaskGenerator(sam)27 28generator = pipeline(model="facebook/sam-vit-base", task="mask-generation", points_per_batch=256)29#image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"30 31controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(32 "SAMControlNet/sd-controlnet-sam-seg", dtype=jnp.float3233)34 35pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(36 "runwayml/stable-diffusion-v1-5",37 controlnet=controlnet,38 revision="flax",39 dtype=jnp.bfloat16,40)41 42params["controlnet"] = controlnet_params43p_params = replicate(params)44 45 46with gr.Blocks() as demo:47 gr.Markdown("# WildSynth: Synthetic Wildlife Data Generation")48 gr.Markdown(49 """50 ## Work in Progress51 ### About52 We have trained a JAX ControlNet model for semantic segmentation on Wildlife Animal Images.53 54 For the training data creation we used the [Wildlife Animals Images](https://www.kaggle.com/datasets/anshulmehtakaggl/wildlife-animals-images) dataset.55 We created segmentation masks with the help of [Grounded SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything) where we used the animals names 56 as input prompts for detection and more accurate segmentation.57 58 ### How To Use59 60 """61 )62 with gr.Row():63 input_img = gr.Image(label="Input", type="pil")64 mask_img = gr.Image(label="Mask", interactive=False)65 output_img = gr.Image(label="Output", interactive=False)66 67 with gr.Row():68 prompt_text = gr.Textbox(lines=1, label="Prompt")69 negative_prompt_text = gr.Textbox(lines=1, label="Negative Prompt")70 71 with gr.Row():72 submit = gr.Button("Submit")73 clear = gr.Button("Clear")74 75 def generate_mask(image):76 outputs = generator(image, points_per_batch=256)77 mask_images = []78 for mask in outputs["masks"]:79 color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)80 h, w = mask.shape[-2:]81 mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)82 mask_images.append(mask_image)83 84 return np.stack(mask_images)85 86 # predictor.set_image(image)87 # input_point = np.array([120, 21])88 # input_label = np.ones(input_point.shape[0])89 # mask, _, _ = predictor.predict(90 # point_coords=input_point,91 # point_labels=input_label,92 # multimask_output=False,93 # )94 95 # clear torch cache96 # torch.cuda.empty_cache()97 # mask = Image.fromarray(mask[0, :, :])98 # segs = mask_generator.generate(image)99 # boolean_masks = [s["segmentation"] for s in segs]100 # finseg = np.zeros(101 # (boolean_masks[0].shape[0], boolean_masks[0].shape[1], 3), dtype=np.uint8102 # )103 # # Loop over the boolean masks and assign a unique color to each class104 # for class_id, boolean_mask in enumerate(boolean_masks):105 # hue = class_id * 1.0 / len(boolean_masks)106 # rgb = tuple(int(i * 255) for i in colorsys.hsv_to_rgb(hue, 1, 1))107 # rgb_mask = np.zeros(108 # (boolean_mask.shape[0], boolean_mask.shape[1], 3), dtype=np.uint8109 # )110 # rgb_mask[:, :, 0] = boolean_mask * rgb[0]111 # rgb_mask[:, :, 1] = boolean_mask * rgb[1]112 # rgb_mask[:, :, 2] = boolean_mask * rgb[2]113 # finseg += rgb_mask114 115 # torch.cuda.empty_cache()116 117 # return mask118 119 def infer(120 image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4121 ):122 try:123 rng = jax.random.PRNGKey(int(seed))124 num_inference_steps = int(num_inference_steps)125 image = Image.fromarray(image, mode="RGB")126 num_samples = max(jax.device_count(), int(num_samples))127 p_rng = jax.random.split(rng, jax.device_count())128 129 prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)130 negative_prompt_ids = pipe.prepare_text_inputs(131 [negative_prompts] * num_samples132 )133 processed_image = pipe.prepare_image_inputs([image] * num_samples)134 135 prompt_ids = shard(prompt_ids)136 negative_prompt_ids = shard(negative_prompt_ids)137 processed_image = shard(processed_image)138 139 output = pipe(140 prompt_ids=prompt_ids,141 image=processed_image,142 params=p_params,143 prng_seed=p_rng,144 num_inference_steps=num_inference_steps,145 neg_prompt_ids=negative_prompt_ids,146 jit=True,147 ).images148 149 del negative_prompt_ids150 del processed_image151 del prompt_ids152 153 output = output.reshape((num_samples,) + output.shape[-3:])154 final_image = [np.array(x * 255, dtype=np.uint8) for x in output]155 print(output.shape)156 del output157 158 except Exception as e:159 print("Error: " + str(e))160 final_image = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples161 finally:162 gc.collect()163 return final_image164 165 def _clear(sel_pix, img, mask, seg, out, prompt, neg_prompt, bg):166 img = None167 mask = None168 seg = None169 out = None170 prompt = ""171 neg_prompt = ""172 bg = False173 return img, mask, seg, out, prompt, neg_prompt, bg174 175 input_img.change(176 generate_mask,177 inputs=[input_img],178 outputs=[mask_img],179 )180 submit.click(181 infer,182 inputs=[mask_img, prompt_text, negative_prompt_text],183 outputs=[output_img],184 )185 clear.click(186 _clear,187 inputs=[188 input_img,189 mask_img,190 output_img,191 prompt_text,192 negative_prompt_text,193 ],194 outputs=[195 input_img,196 mask_img,197 output_img,198 prompt_text,199 negative_prompt_text,200 ],201 )202 203if __name__ == "__main__":204 demo.queue()205 demo.launch()206