Photo Mosaic App on Pythonista (iOS app)

mosaic_to_photo.py

import appex
import sys
import photos
import ui
import io
from PIL import Image, ImageOps, ImageDraw, ImageFilter, ImageStat
import random
import pathlib

def scale_image(pil_img, max_dim=1024):
    """
    Scales down pil_img so the longest side is at most max_dim.
    Returns (scaled_image, scale_factor).
    scale_factor = original_size / new_size
    """
    w, h = pil_img.size
    scale = 1.0
    if max(w, h) > max_dim:
        scale = max(w, h) / max_dim
        new_w = int(w / scale)
        new_h = int(h / scale)
        scaled = pil_img.resize((new_w, new_h), Image.BILINEAR)
        return scaled, scale
    else:
        return pil_img.copy(), 1.0

def apply_mosaic(pil_img, rect, block_size=30):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    small = region.resize((max(1, w // block_size),
                           max(1, h // block_size)),
                          resample=Image.NEAREST)
    mosaic_region = small.resize(region.size, Image.NEAREST)
    pil_img.paste(mosaic_region, (x0, y0, x1, y1))

def apply_blur(pil_img, rect, radius=10):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    blurred = region.filter(ImageFilter.GaussianBlur(radius))
    pil_img.paste(blurred, (x0, y0, x1, y1))

def apply_circular_mosaic(pil_img, rect, block_size=30):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    small = region.resize((max(1, w // block_size),
                           max(1, h // block_size)),
                          resample=Image.NEAREST)
    mosaic_region = small.resize(region.size, Image.NEAREST)
    mask = Image.new("L", (w, h), 0)
    draw = ImageDraw.Draw(mask)
    draw.ellipse((0, 0, w, h), fill=255)
    region.paste(mosaic_region, (0, 0), mask)
    pil_img.paste(region, (x0, y0, x1, y1))

def apply_circular_blur(pil_img, rect, radius=10):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    blurred = region.filter(ImageFilter.GaussianBlur(radius))
    mask = Image.new("L", (w, h), 0)
    draw = ImageDraw.Draw(mask)
    draw.ellipse((0, 0, w, h), fill=255)
    pil_img.paste(blurred, (x0, y0, x1, y1), mask)

def apply_marble(pil_img, rect, block_size=30):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    width, height = region.size
    output_region = Image.new('RGB', (width, height), (255, 255, 255))
    draw = ImageDraw.Draw(output_region)
    radius = block_size // 2
    grid_size = radius * 2
    for gy in range(0, height, grid_size):
        for gx in range(0, width, grid_size):
            offset_x = random.randint(-radius//2, radius//2)
            offset_y = random.randint(-radius//2, radius//2)
            center_x = gx + radius + offset_x
            center_y = gy + radius + offset_y
            box = (gx, gy, min(gx + grid_size, width), min(gy + grid_size, height))
            sample_region = region.crop(box)
            stat = ImageStat.Stat(sample_region)
            r, g, b = stat.mean[:3]
            bounding_box = [
                (center_x - radius, center_y - radius),
                (center_x + radius, center_y + radius)
            ]
            draw.ellipse(bounding_box, fill=(int(r), int(g), int(b)))
    pil_img.paste(output_region, (x0, y0, x1, y1))

def apply_circular_marble(pil_img, rect, block_size=30):
    x0, y0, w, h = map(int, rect)
    if w <= 0 or h <= 0:
        return
    x1, y1 = x0 + w, y0 + h
    region = pil_img.crop((x0, y0, x1, y1))
    width, height = region.size
    output_region = Image.new('RGB', (width, height), (255, 255, 255))
    draw = ImageDraw.Draw(output_region)
    radius = block_size // 2
    grid_size = radius * 2
    for gy in range(0, height, grid_size):
        for gx in range(0, width, grid_size):
            offset_x = random.randint(-radius//2, radius//2)
            offset_y = random.randint(-radius//2, radius//2)
            center_x = gx + radius + offset_x
            center_y = gy + radius + offset_y
            box = (gx, gy, min(gx + grid_size, width), min(gy + grid_size, height))
            sample_region = region.crop(box)
            stat = ImageStat.Stat(sample_region)
            r, g, b = stat.mean[:3]
            bounding_box = [
                (center_x - radius, center_y - radius),
                (center_x + radius, center_y + radius)
            ]
            draw.ellipse(bounding_box, fill=(int(r), int(g), int(b)))
    mask = Image.new("L", (w, h), 0)
    draw_mask = ImageDraw.Draw(mask)
    draw_mask.ellipse((0, 0, w, h), fill=255)
    pil_img.paste(output_region, (x0, y0, x1, y1), mask)

def adjust_rect_for_front(rect, img_size):
    x, y, w, h = rect
    fw, fh = img_size
    return (fw - (x + w), fh - (y + h), w, h)

class OverlayView(ui.View):
    def __init__(self, mosaic_view, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mosaic_view = mosaic_view
        self.background_color = None
        self.user_interaction_enabled = False

    def draw(self):
        current_rect = self.mosaic_view.current_rect
        if current_rect:
            ui.set_color((1, 0, 0, 0.3))
            iv = self.mosaic_view.img_view
            rel_rect = (
                current_rect[0] - iv.x,
                current_rect[1] - iv.y,
                current_rect[2],
                current_rect[3]
            )
            # Check if circular shape is selected
            if self.mosaic_view.shape_selector.selected_index == 1:
                # Draw circular overlay
                path = ui.Path.oval(*rel_rect)
            else:
                # Draw rectangular overlay
                path = ui.Path.rect(*rel_rect)
            path.line_width = 2
            path.stroke()

class MosaicView(ui.View):
    def __init__(self, pil_img, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.background_color = 'white'
        pil_img = ImageOps.exif_transpose(pil_img)
        self.fullres_img = pil_img
        self.small_original, self.scale_factor = scale_image(pil_img, max_dim=1024)
        self.small_edited = self.small_original.copy()
        self.is_front = False
        self.selections = []
        self.block_size = 30

        self.slider = ui.Slider()
        self.slider.continuous = True
        self.slider.value = (self.block_size - 10) / 70.0
        self.slider.action = self.slider_action
        self.add_subview(self.slider)

        self.effect_selector = ui.SegmentedControl()
        self.effect_selector.segments = ["Mosaic", "Blur", "Marble"]
        self.effect_selector.selected_index = 0
        self.add_subview(self.effect_selector)

        self.shape_selector = ui.SegmentedControl()
        self.shape_selector.segments = ["Rectangular", "Circular"]
        self.shape_selector.selected_index = 0
        self.add_subview(self.shape_selector)

        self.img_view = ui.ImageView()
        self.img_view.content_mode = ui.CONTENT_SCALE_ASPECT_FIT
        self.add_subview(self.img_view)

        self.overlay = OverlayView(self)
        self.img_view.add_subview(self.overlay)

        self.start_point = None
        self.current_rect = None

        done_btn = ui.ButtonItem(title='Done', action=self.done_action)
        undo_btn = ui.ButtonItem(title='Undo', action=self.undo_action)
        self.left_button_items = [undo_btn]
        self.right_button_items = [done_btn]

        self.update_image()

    def layout(self):
        self.slider.frame = (0, 0, self.width, 40)
        self.effect_selector.frame = (0, 40, self.width, 32)
        self.shape_selector.frame = (0, 72, self.width, 32)
        top_margin = 104

        img_width = self.width * 0.9
        img_x = (self.width - img_width) / 2
        img_height = self.height - top_margin

        self.img_view.frame = (img_x, top_margin, img_width, img_height)
        self.overlay.frame = (0, 0, self.img_view.width, self.img_view.height)

    def slider_action(self, sender):
        self.block_size = int(10 + sender.value * 70)
        self.reapply_all_selections()

    def reapply_all_selections(self):
        self.small_edited = self.small_original.copy()
        for (rect, _, eff) in self.selections:
            adj_rect = rect
            if eff == 0:
                apply_mosaic(self.small_edited, adj_rect, block_size=self.block_size)
            elif eff == 1:
                apply_blur(self.small_edited, adj_rect, radius=self.block_size)
            elif eff == 2:
                apply_circular_mosaic(self.small_edited, adj_rect, block_size=self.block_size)
            elif eff == 3:
                apply_circular_blur(self.small_edited, adj_rect, radius=self.block_size)
            elif eff == 4:
                apply_marble(self.small_edited, adj_rect, block_size=self.block_size)
            elif eff == 5:
                apply_circular_marble(self.small_edited, adj_rect, block_size=self.block_size)
        self.update_image()

    def update_image(self):
        buf = io.BytesIO()
        self.small_edited.save(buf, format='PNG')
        buf.seek(0)
        self.img_view.image = ui.Image.from_data(buf.read())
        self.set_needs_display()
        self.overlay.set_needs_display()

    def touch_began(self, touch):
        self.start_point = touch.location
        self.current_rect = None

    def touch_moved(self, touch):
        end_point = touch.location
        x0, y0 = self.start_point
        x1, y1 = end_point
        self.current_rect = (min(x0, x1),
                             min(y0, y1),
                             abs(x1 - x0),
                             abs(y1 - y0))
        self.set_needs_display()
        self.overlay.set_needs_display()

    def touch_ended(self, touch):
        if not self.start_point or not self.current_rect:
            return

        sp = self.view_to_small_coords(self.start_point)
        ep = self.view_to_small_coords(touch.location)

        x0, y0 = sp
        x1, y1 = ep
        w, h = abs(x1 - x0), abs(y1 - y0)
        effect_rect = (min(x0, x1), min(y0, y1), w, h)

        effect = self.effect_selector.selected_index
        shape = self.shape_selector.selected_index

        if effect == 0 and shape == 0:
            current_effect = 0
        elif effect == 1 and shape == 0:
            current_effect = 1
        elif effect == 0 and shape == 1:
            current_effect = 2
        elif effect == 1 and shape == 1:
            current_effect = 3
        elif effect == 2 and shape == 0:
            current_effect = 4
        elif effect == 2 and shape == 1:
            current_effect = 5

        adj_rect = effect_rect

        if current_effect == 0:
            apply_mosaic(self.small_edited, adj_rect, block_size=self.block_size)
        elif current_effect == 1:
            apply_blur(self.small_edited, adj_rect, radius=self.block_size)
        elif current_effect == 2:
            apply_circular_mosaic(self.small_edited, adj_rect, block_size=self.block_size)
        elif current_effect == 3:
            apply_circular_blur(self.small_edited, adj_rect, radius=self.block_size)
        elif current_effect == 4:
            apply_marble(self.small_edited, adj_rect, block_size=self.block_size)
        elif current_effect == 5:
            apply_circular_marble(self.small_edited, adj_rect, block_size=self.block_size)

        self.selections.append((effect_rect, self.block_size, current_effect))
        self.current_rect = None
        self.update_image()

    def view_to_small_coords(self, pt):
        vx, vy = pt
        vw, vh = self.img_view.width, self.img_view.height
        iw, ih = self.small_edited.size
        scale = min(vw / iw, vh / ih)
        offset_x = (vw - (iw * scale)) / 2
        offset_y = (vh - (ih * scale)) / 2
        img_x = (vx - self.img_view.x - offset_x) / scale
        img_y = (vy - self.img_view.y - offset_y) / scale
        return (max(0, min(iw, img_x)), max(0, min(ih, img_y)))

    def undo_action(self, sender):
        if self.selections:
            self.selections.pop()
            self.reapply_all_selections()

    def done_action(self, sender):
        final_img = self.fullres_img.copy()
        for (small_rect, _, eff) in self.selections:
            sx, sy, sw, sh = small_rect
            x = int(sx * self.scale_factor)
            y = int(sy * self.scale_factor)
            w = int(sw * self.scale_factor)
            h = int(sh * self.scale_factor)
            full_rect = (x, y, w, h)
            full_block_size = int(self.block_size * self.scale_factor)
            if eff == 0:
                apply_mosaic(final_img, full_rect, block_size=full_block_size)
            elif eff == 1:
                apply_blur(final_img, full_rect, radius=full_block_size)
            elif eff == 2:
                apply_circular_mosaic(final_img, full_rect, block_size=full_block_size)
            elif eff == 3:
                apply_circular_blur(final_img, full_rect, radius=full_block_size)
            elif eff == 4:
                apply_marble(final_img, full_rect, block_size=full_block_size)
            elif eff == 5:
                apply_circular_marble(final_img, full_rect, block_size=full_block_size)

        buf = io.BytesIO()
        final_img.save(buf, format='JPEG', quality=90, icc_profile=final_img.info.get('icc_profile'))
        buf.seek(0)
        ui_img = ui.Image.from_data(buf.read())
        photos.save_image(ui_img)
        self.close()

def main():
    # Check if a file path argument is provided
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
        try:
            pil_img = Image.open(file_path)
        except Exception as e:
            print("Failed to open image from argument:", e)
            return
    elif appex.is_running_extension():
        # In extension mode, save the image and open the full app
        pil_img = appex.get_image('pil')
        if not pil_img:
            print("No image found in extension input.")
            return
        tmp_path = pathlib.Path().home().joinpath('Documents/tmp_image.jpg')
        pil_img.save(tmp_path, format='JPEG', quality=95, icc_profile=pil_img.info.get('icc_profile'), exif=pil_img.info.get('exif'))
        # Open this script in full Pythonista app with the image path as argument
        url = 'pythonista3://mosaic_to_photo?action=run&root=icloud&argv=' + tmp_path.as_posix()
        import webbrowser
        webbrowser.open(url)
        return
    else:
        pil_img = photos.pick_image(original=True)
        if not pil_img:
            print("No image selected.")
            return

    v = MosaicView(pil_img, frame=(0, 0, 0, 0))
    v.present(style='fullscreen', orientations=['portrait'])

if __name__ == '__main__':
    main()