Upload 3 files

engine.py

@@ -0,0 +1,97 @@
+import torch
+from torchvision import transforms
+from torchvision.transforms.functional import to_pil_image, to_tensor
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
+from diffusers import StableDiffusionInpaintPipeline, EulerDiscreteScheduler
+from PIL import Image
+import numpy as np
+import cv2
+
+class SegmentAnythingModel:
+    def __init__(self, sam_checkpoint, model_type, device):
+        self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+        self.sam.to(device=device)
+        self.mask_generator = SamAutomaticMaskGenerator(
+            model=self.sam,
+            points_per_side=32,
+            pred_iou_thresh=0.99,
+            stability_score_thresh=0.92,
+            crop_n_layers=1,
+            crop_n_points_downscale_factor=2,
+            min_mask_region_area=100
+        )
+        self.target_size = (512, 512)
+
+    def preprocess_image(self, image):
+        """Resize image to 512x512"""
+        if isinstance(image, str):
+            image = Image.open(image)
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        # Get current dimensions
+        width, height = image.size
+        
+        # Resize to 512x512 directly
+        image = image.resize(self.target_size, Image.Resampling.LANCZOS)
+        return np.array(image)
+
+    def generate_masks(self, image):
+        processed_image = self.preprocess_image(image)
+        return self.mask_generator.generate(processed_image)
+
+class StableDiffusionInpaintingPipeline:
+    def __init__(self, model_dir):
+        # Initialize the scheduler first
+        self.scheduler = EulerDiscreteScheduler.from_pretrained(model_dir, subfolder="scheduler")
+        
+        # Initialize the pipeline with the scheduler
+        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            model_dir,
+            scheduler=self.scheduler,
+            revision="fp16",
+            torch_dtype=torch.float16
+        )
+        self.pipe = self.pipe.to("cuda")
+        self.pipe.enable_xformers_memory_efficient_attention()
+        self.target_size = (512, 512)
+
+    def preprocess_image(self, image):
+        """Ensure image is in the right format and size"""
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        return image.resize(self.target_size, Image.Resampling.LANCZOS)
+
+    def inpaint(self, prompt, image, mask_image, guidance_scale=10, num_inference_steps=60, generator=None):
+        """
+        Args:
+            prompt (str): The prompt for inpainting
+            image (PIL.Image or np.ndarray): The original image
+            mask_image (PIL.Image or np.ndarray): The mask for inpainting
+            guidance_scale (float): Higher guidance scale encourages images that are closer to the prompt
+            num_inference_steps (int): Number of denoising steps
+            generator (torch.Generator): Generator for reproducibility
+        """
+        # Preprocess images
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        if isinstance(mask_image, np.ndarray):
+            mask_image = Image.fromarray(mask_image)
+
+        # Resize images
+        image = image.resize(self.target_size, Image.Resampling.LANCZOS)
+        mask_image = mask_image.resize(self.target_size, Image.Resampling.NEAREST)
+
+        # Run inpainting
+        output = self.pipe(
+            prompt=prompt,
+            image=image,
+            mask_image=mask_image,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            height=512,
+            width=512
+        )
+        
+        return output.images[0]

requirements.txt

@@ -0,0 +1,18 @@
+regex
+tqdm
+diffusers
+transformers
+scipy
+accelerate
+opencv-python
+Xformers
+gradio
+torch
+torchvision
+Pillow
+matplotlib
+numpy
+git+https://github.com/facebookresearch/segment-anything.git
+pycocotools 
+onnxruntime 
+onnx

utils.py

@@ -0,0 +1,97 @@
+import PIL
+import numpy as np
+import copy
+import cv2
+import matplotlib.pyplot as plt
+from torchvision.transforms.functional import to_pil_image
+import torch
+from PIL import Image
+import matplotlib
+matplotlib.use('Agg')
+
+def show_anns(anns, ax=None):
+    if len(anns) == 0:
+        return
+    if ax is None:
+        ax = plt.gca()
+
+    sorted_anns = sorted(enumerate(anns), key=(lambda x: x[1]['area']), reverse=True)
+
+    for original_idx, ann in sorted_anns:
+        m = ann['segmentation']
+        if m.shape != (512, 512):  # Ensure mask is right size
+            m = cv2.resize(m.astype(float), (512, 512))
+        
+        # Create a random color for this mask
+        color_mask = np.random.random(3)
+        
+        # Create the colored mask
+        colored_mask = np.zeros((512, 512, 3))
+        for i in range(3):
+            colored_mask[:,:,i] = color_mask[i]
+        
+        # Add the mask with transparency
+        ax.imshow(np.dstack([colored_mask, m * 0.35]))
+        
+        # Find contours of the mask
+        contours, _ = cv2.findContours((m * 255).astype(np.uint8), 
+                                     cv2.RETR_EXTERNAL, 
+                                     cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Add mask number if contours exist
+        if contours:
+            # Get the largest contour
+            cnt = max(contours, key=cv2.contourArea)
+            M = cv2.moments(cnt)
+            
+            if M["m00"] != 0:
+                cx = int(M["m10"] / M["m00"])
+                cy = int(M["m01"] / M["m00"])
+                
+                # Add text with white color and black outline for visibility
+                ax.text(cx, cy, str(original_idx), 
+                       color='white', 
+                       fontsize=16, 
+                       ha='center', 
+                       va='center', 
+                       fontweight='bold',
+                       bbox=dict(facecolor='black', 
+                                alpha=0.5, 
+                                edgecolor='none', 
+                                pad=1))
+
+
+def create_image_grid(original_image, images, names, rows, columns):
+    names = copy.copy(names)
+    images = copy.copy(images)
+    
+    # Filter out empty prompts and their corresponding images
+    filtered_images = []
+    filtered_names = []
+    for img, name in zip(images, names):
+        if name.strip():
+            filtered_images.append(img)
+            filtered_names.append(name)
+    
+    images = filtered_images
+    names = filtered_names
+
+    # Add original image
+    images.insert(0, original_image)
+    names.insert(0, 'Original')
+
+    fig = plt.figure(figsize=(20, 20))
+    
+    for idx, (img, name) in enumerate(zip(images, names)):
+        ax = fig.add_subplot(rows, columns, idx + 1)
+        
+        if isinstance(img, PIL.Image.Image):
+            ax.imshow(img)
+        else:
+            ax.imshow(img)
+            
+        ax.set_title(name, fontsize=12, pad=10)
+        ax.axis('off')
+
+    plt.tight_layout()
+    return fig