Unifying InPaintingGenerator and SyntheticDataGenerator

bitmind/synthetic_data_generation/__init__.py

@@ -1,3 +1,2 @@
 from .synthetic_data_generator import SyntheticDataGenerator
-from .in_painting_generator import InPaintingGenerator
 from .prompt_generator import PromptGenerator

bitmind/synthetic_data_generation/image_utils.py

@@ -1,6 +1,9 @@
+import numpy as np
 import PIL
 import os
-import json
+from PIL import Image, ImageDraw
+from typing import Tuple
+
 from bitmind.validator.config import TARGET_IMAGE_SIZE
 
 
@@ -56,4 +59,51 @@ def save_images_to_disk(image_dataset, start_index, num_images, save_directory,
             image.save(file_path, 'JPEG')  # Save the image
             print(f"Saved: {file_path}")
         except Exception as e:
-            print(f"Failed to save image {i}: {e}")
+            print(f"Failed to save image {i}: {e}")
+
+
+def create_random_mask(size: Tuple[int, int]) -> Image.Image:
+    """
+    Create a random mask for i2i transformation.
+    """
+    w, h = size
+    mask = Image.new('RGB', size, 'black')
+
+    if np.random.rand() < 0.5:
+        # Rectangular mask with smoother edges
+        width = np.random.randint(w//4, w//2)
+        height = np.random.randint(h//4, h//2)
+
+        # Center the rectangle with some random offset
+        x1 = (w - width) // 2 + np.random.randint(-width//4, width//4)
+        y1 = (h - height) // 2 + np.random.randint(-height//4, height//4)
+
+        # Create mask with PIL draw for smoother edges
+        draw = ImageDraw.Draw(mask)
+        draw.rounded_rectangle(
+            [x1, y1, x1 + width, y1 + height],
+            radius=min(width, height) // 10,  # Smooth corners
+            fill='white'
+        )
+    else:
+        # Circular mask with feathered edges
+        draw = ImageDraw.Draw(mask)
+        center_x = w//2
+        center_y = h//2
+
+        # Make radius proportional to image size
+        radius = min(w, h) // 4
+
+        # Add small random offset to center
+        center_x += np.random.randint(-radius//4, radius//4)
+        center_y += np.random.randint(-radius//4, radius//4)
+
+        # Draw multiple circles with decreasing opacity for feathered edge
+        for r in range(radius, radius-10, -1):
+            opacity = int(255 * (r - (radius-10)) / 10)
+            draw.ellipse(
+                [center_x-r, center_y-r, center_x+r, center_y+r],
+                fill=(255, 255, 255, opacity)
+            )
+
+    return mask