try subplotting

app.py

@@ -1,7 +1,172 @@
 import gradio as gr
+import numpy as np
+import torch
+import jax
+import jax.numpy as jnp
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+from PIL import Image
+from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
+from diffusers import (
+    FlaxStableDiffusionControlNetPipeline,
+    FlaxControlNetModel,
+)
+from transformers import pipeline
 
-def greet(name):
-    return "Hello " + name + "!!"
+import colorsys
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+sam_checkpoint = "sam_vit_h_4b8939.pth"
+model_type = "vit_h"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+#sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+#sam.to(device=device)
+#predictor = SamPredictor(sam)
+#mask_generator = SamAutomaticMaskGenerator(sam)
+
+generator = pipeline(model="facebook/sam-vit-base", task="mask-generation", points_per_batch=256)
+#image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
+
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "SAMControlNet/sd-controlnet-sam-seg", dtype=jnp.float32
+)
+
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    revision="flax",
+    dtype=jnp.bfloat16,
+)
+
+params["controlnet"] = controlnet_params
+p_params = replicate(params)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Ahsans version WildSynth: Synthetic Wildlife Data Generation")
+    gr.Markdown(
+        """
+        ## Work in Progress
+        ### About
+        We have trained a JAX ControlNet model for semantic segmentation on Wildlife Animal Images.
+        
+        For the training data creation we used the [Wildlife Animals Images](https://www.kaggle.com/datasets/anshulmehtakaggl/wildlife-animals-images) dataset.
+        We created segmentation masks with the help of [Grounded SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything) where we used the animals names 
+        as input prompts for detection and more accurate segmentation.
+        
+        ### How To Use
+        
+    """
+    )
+    with gr.Row():
+        input_img = gr.Image(label="Input", type="pil")
+        mask_img = gr.Image(label="Mask", interactive=False)
+        output_img = gr.Image(label="Output", interactive=False)
+
+    with gr.Row():
+        prompt_text = gr.Textbox(lines=1, label="Prompt")
+        negative_prompt_text = gr.Textbox(lines=1, label="Negative Prompt")
+
+    with gr.Row():
+        submit = gr.Button("Submit")
+        clear = gr.Button("Clear")
+
+    def generate_mask(image):
+        outputs = generator(image, points_per_batch=256)
+        mask_images = []
+        for mask in outputs["masks"]:
+            color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)
+            h, w = mask.shape[-2:]
+            mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            mask_images.append(mask_image)
+        
+        return np.stack(mask_images)
+
+    def infer(
+        image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4
+    ):
+        try:
+            rng = jax.random.PRNGKey(int(seed))
+            num_inference_steps = int(num_inference_steps)
+            image = Image.fromarray(image, mode="RGB")
+            num_samples = max(jax.device_count(), int(num_samples))
+            p_rng = jax.random.split(rng, jax.device_count())
+
+            prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
+            negative_prompt_ids = pipe.prepare_text_inputs(
+                [negative_prompts] * num_samples
+            )
+            processed_image = pipe.prepare_image_inputs([image] * num_samples)
+
+            prompt_ids = shard(prompt_ids)
+            negative_prompt_ids = shard(negative_prompt_ids)
+            processed_image = shard(processed_image)
+
+            output = pipe(
+                prompt_ids=prompt_ids,
+                image=processed_image,
+                params=p_params,
+                prng_seed=p_rng,
+                num_inference_steps=num_inference_steps,
+                neg_prompt_ids=negative_prompt_ids,
+                jit=True,
+            ).images
+
+            del negative_prompt_ids
+            del processed_image
+            del prompt_ids
+
+            output = output.reshape((num_samples,) + output.shape[-3:])
+            final_image = [np.array(x * 255, dtype=np.uint8) for x in output]
+            print(output.shape)
+            del output
+
+        except Exception as e:
+            print("Error: " + str(e))
+            final_image = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples
+        finally:
+            gc.collect()
+            return final_image
+
+    def _clear(sel_pix, img, mask, seg, out, prompt, neg_prompt, bg):
+        img = None
+        mask = None
+        seg = None
+        out = None
+        prompt = ""
+        neg_prompt = ""
+        bg = False
+        return img, mask, seg, out, prompt, neg_prompt, bg
+
+    input_img.change(
+        generate_mask,
+        inputs=[input_img],
+        outputs=[mask_img],
+    )
+    submit.click(
+        infer,
+        inputs=[mask_img, prompt_text, negative_prompt_text],
+        outputs=[output_img],
+    )
+    clear.click(
+        _clear,
+        inputs=[
+            input_img,
+            mask_img,
+            output_img,
+            prompt_text,
+            negative_prompt_text,
+        ],
+        outputs=[
+            input_img,
+            mask_img,
+            output_img,
+            prompt_text,
+            negative_prompt_text,
+        ],
+    )
+
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch()

app.py.orig

@@ -0,0 +1,122 @@
+import gradio as gr
+
+import torch
+from PIL import Image
+import requests
+from transformers import SamModel, SamProcessor
+import numpy as np
+
+
+def show_mask(mask, ax, random_color=False):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        color = np.array([30/255, 144/255, 255/255, 0.6])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    ax.imshow(mask_image)
+
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))  
+
+def show_boxes_on_image(raw_image, boxes):
+    plt.figure(figsize=(10,10))
+    plt.imshow(raw_image)
+    for box in boxes:
+      show_box(box, plt.gca())
+    plt.axis('on')
+    plt.show()
+
+
+def show_points_on_image(raw_image, input_points, input_labels=None):
+    plt.figure(figsize=(10,10))
+    plt.imshow(raw_image)
+    input_points = np.array(input_points)
+    if input_labels is None:
+      labels = np.ones_like(input_points[:, 0])
+    else:
+      labels = np.array(input_labels)
+    show_points(input_points, labels, plt.gca())
+    plt.axis('on')
+    plt.show()
+
+def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None):
+    plt.figure(figsize=(10,10))
+    plt.imshow(raw_image)
+    input_points = np.array(input_points)
+    if input_labels is None:
+      labels = np.ones_like(input_points[:, 0])
+    else:
+      labels = np.array(input_labels)
+    show_points(input_points, labels, plt.gca())
+    for box in boxes:
+      show_box(box, plt.gca())
+    plt.axis('on')
+    plt.show()
+
+
+def show_points_and_boxes_on_image(raw_image, boxes, input_points, input_labels=None):
+    plt.figure(figsize=(10,10))
+    plt.imshow(raw_image)
+    input_points = np.array(input_points)
+    if input_labels is None:
+      labels = np.ones_like(input_points[:, 0])
+    else:
+      labels = np.array(input_labels)
+    show_points(input_points, labels, plt.gca())
+    for box in boxes:
+      show_box(box, plt.gca())
+    plt.axis('on')
+    plt.show()
+
+
+def show_points(coords, labels, ax, marker_size=375):
+    pos_points = coords[labels==1]
+    neg_points = coords[labels==0]
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
+
+
+def apply_masks_on_image(raw_image, masks, scores):
+    if len(masks.shape) == 4:
+      masks = masks.squeeze()
+    if scores.shape[0] == 1:
+      scores = scores.squeeze()
+
+    nb_predictions = scores.shape[-1]
+    fig, axes = plt.subplots(1, nb_predictions, figsize=(15, 15))
+
+    for i, (mask, score) in enumerate(zip(masks, scores)):
+      mask = mask.cpu().detach()
+      axes[i].imshow(np.array(raw_image))
+      show_mask(mask, axes[i])
+      axes[i].title.set_text(f"Mask {i+1}, Score: {score.item():.3f}")
+      axes[i].axis("off")
+    plt.show()
+
+def segment(imageUrl):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+    processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+
+    img_url = imageUrl#"https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
+    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
+    input_points = [[[450, 600]]]  # 2D location of a window in the image
+
+    inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+    )
+    scores = outputs.iou_scores
+    return {"Masks": masks, "Scores": scores}
+
+gr.Interface(fn=predict, 
+             inputs=gr.Image(type="pil"),
+             outputs=[{"type":"dataframe","name":"Categories Scores"},
+             {"type":"dataframe","name":"Categories Labels"}],
+  ).launch()