optiminization

python/transition.py

@@ -6,10 +6,11 @@ from matplotlib.tri import Triangulation
 
 from image_preparing import fix_image_resolution
 
+import time
+
+
+transformation_resolution = (1600, 1200)
 
-def prepare_image(path):
-    pil = Image.open(path).convert("RGB").resize((800, 600))
-    return pil
 def auto_correspondences_grid(pil, nx=8, ny=6):
     w, h = pil.size
     xs = np.linspace(0, w - 1, nx)
@@ -68,89 +69,111 @@ def reconstruct_goal_from_source(src_pil, pts1, pts2):
         out.paste(warped, (min_x, min_y), mask)
     return out
 def transform_images(src_path, goal_path):
-    fix_image_resolution(src_path)
-    fix_image_resolution(goal_path)
-    src = prepare_image(src_path)
-    goal = prepare_image(goal_path)
-    frames = tile_reposition_transition(src, goal, grid=(40, 30), n_frames=24)
-    fig, ax = plt.subplots(figsize=(8, 6))
-    im = ax.imshow(frames[0])
+    src = fix_image_resolution(src_path, size=transformation_resolution)
+    goal = fix_image_resolution(goal_path, size=transformation_resolution)
+    start = time.time()
+    image = tile_reposition_transition(src, goal, grid=transformation_resolution)
+    end = time.time()
+    print(f"Transition took {end - start:.2f} seconds")
+    fig, ax = plt.subplots(figsize=(8, 6), facecolor="black")
+    fig.patch.set_facecolor("black")
+    ax.set_facecolor("black")
+    ax.imshow(image, cmap=None)
     ax.axis("off")
-    for frame in frames:
-        im.set_data(frame)
-        fig.canvas.draw_idle()
-        plt.pause(0.05)
+
+    def _close_on_key(event):
+        if event.key in ("q", "escape"):
+            plt.close(fig)
+
+    fig.canvas.mpl_connect("key_press_event", _close_on_key)
+
     plt.show()
-def tile_reposition_transition(src_pil, goal_pil, grid=(40, 30), n_frames=24):
-    
-    src = src_pil.convert("RGBA")
-    goal = goal_pil.convert("RGBA")
+def _grid_edges(size, divisions):
+    return np.rint(np.linspace(0, size, divisions + 1)).astype(np.int32)
+def _build_tile_boxes_and_centers(w, h, nx, ny):
+    x_edges = _grid_edges(w, nx)
+    y_edges = _grid_edges(h, ny)
+
+    x0 = np.repeat(x_edges[:-1], ny)
+    x1 = np.repeat(x_edges[1:], ny)
+    y0 = np.tile(y_edges[:-1], nx)
+    y1 = np.tile(y_edges[1:], nx)
+
+    boxes = np.stack([x0, y0, x1, y1], axis=1)
+    centers = np.stack([(x0 + x1) * 0.5, (y0 + y1) * 0.5], axis=1).astype(np.float32)
+    areas = np.maximum((x1 - x0) * (y1 - y0), 1).astype(np.float32)
+    return boxes, centers, areas
+def _integral_image_rgb(arr):
+    integral = arr.cumsum(axis=0).cumsum(axis=1)
+    return np.pad(integral, ((1, 0), (1, 0), (0, 0)), mode="constant", constant_values=0)
+def _tile_means_from_integral(integral, boxes, areas):
+    x0 = boxes[:, 0]
+    y0 = boxes[:, 1]
+    x1 = boxes[:, 2]
+    y1 = boxes[:, 3]
+
+    sums = (
+        integral[y1, x1]
+        - integral[y0, x1]
+        - integral[y1, x0]
+        + integral[y0, x0]
+    )
+    return sums / areas[:, None]
+def _match_tiles_by_feature_rank(src_avgs, src_centers, goal_avgs, goal_centers, w, h):
+    color_scale = 255.0 * np.sqrt(3.0)
+    spatial_scale = max(np.hypot(w, h), 1e-6)
+
+    src_feat = np.hstack([
+        src_avgs / color_scale,
+        (src_centers / spatial_scale) * 0.35,
+    ])
+    goal_feat = np.hstack([
+        goal_avgs / color_scale,
+        (goal_centers / spatial_scale) * 0.35,
+    ])
+
+    # Fixed projection keeps pairing deterministic and avoids O(n^2) matching.
+    projection = np.array([0.50, 0.30, 0.20, 0.60, 0.40], dtype=np.float32)
+    src_order = np.argsort(src_feat @ projection)
+    goal_order = np.argsort(goal_feat @ projection)
+    return src_order, goal_order
+def tile_reposition_transition(src_pil, goal_pil, grid=transformation_resolution):
+    src = src_pil.convert("RGB")
+    goal = goal_pil.convert("RGB")
     w, h = src.size
     nx, ny = grid
 
-    def make_tiles(pil):
-        tiles = []
-        for iy in range(ny):
-            y0 = int(round(iy * h / ny))
-            y1 = int(round((iy + 1) * h / ny))
-            for ix in range(nx):
-                x0 = int(round(ix * w / nx))
-                x1 = int(round((ix + 1) * w / nx))
-                box = (x0, y0, x1, y1)
-                crop = pil.crop(box).convert("RGB")
-                arr = np.array(crop, dtype=float)
-                avg = arr.reshape(-1, 3).mean(axis=0)
-                center = ((x0 + x1) / 2.0, (y0 + y1) / 2.0)
-                tiles.append({"box": box, "img": crop, "avg": avg, "center": center})
-        return tiles
+    src_boxes, src_centers, src_areas = _build_tile_boxes_and_centers(w, h, nx, ny)
+    goal_boxes, goal_centers, goal_areas = _build_tile_boxes_and_centers(w, h, nx, ny)
 
-    src_tiles = make_tiles(src)
-    goal_tiles = make_tiles(goal)
+    src_arr = np.asarray(src, dtype=np.float32)
+    goal_arr = np.asarray(goal, dtype=np.float32)
+    src_integral = _integral_image_rgb(src_arr)
+    goal_integral = _integral_image_rgb(goal_arr)
 
-    max_spatial = np.hypot(w, h)
-    src_available = set(range(len(src_tiles)))
-    mapping = {} 
-    for g_idx, g in enumerate(goal_tiles):
-        best = None
-        best_cost = None
-        for s_idx in list(src_available):
-            s = src_tiles[s_idx]
-            color_dist = np.linalg.norm(g["avg"] - s["avg"]) / (255.0 * np.sqrt(3))
-            spatial_dist = np.linalg.norm(np.array(g["center"]) - np.array(s["center"])) / max_spatial
-            cost = color_dist + 0.35 * spatial_dist
-            if best_cost is None or cost < best_cost:
-                best_cost = cost
-                best = s_idx
-        if best is None:
-            best = src_available.pop()
-        else:
-            src_available.remove(best)
-        mapping[best] = g_idx
+    src_avgs = _tile_means_from_integral(src_integral, src_boxes, src_areas)
+    goal_avgs = _tile_means_from_integral(goal_integral, goal_boxes, goal_areas)
 
-    targets = {}
-    for s_idx, g_idx in mapping.items():
-        s = src_tiles[s_idx]
-        g = goal_tiles[g_idx]
-        targets[s_idx] = {"start": (s["box"][0], s["box"][1]), "end": (g["box"][0], g["box"][1]), "img": s["img"], "size": (s["box"][2]-s["box"][0], s["box"][3]-s["box"][1])}
-
-    frames = []
-    for k in range(n_frames + 1):
-        t = k / float(n_frames)
-        frame = Image.new("RGBA", (w, h), (0, 0, 0, 255))
-        for s_idx, info in targets.items():
-            sx, sy = info["start"]
-            ex, ey = info["end"]
-            cx = int(round(sx + (ex - sx) * t))
-            cy = int(round(sy + (ey - sy) * t))
-            frame.paste(info["img"], (cx, cy))
-        frames.append(frame.convert("RGB"))
+    src_order, goal_order = _match_tiles_by_feature_rank(
+        src_avgs,
+        src_centers,
+        goal_avgs,
+        goal_centers,
+        w,
+        h,
+    )
 
     final = Image.new("RGB", (w, h))
-    for s_idx, g_idx in mapping.items():
-        src_img = src_tiles[s_idx]["img"].convert("RGB")
-        x0, y0, x1, y1 = goal_tiles[g_idx]["box"]
-        final.paste(src_img, (x0, y0))
-    frames[-1] = final
-    return frames
+    for s_idx, g_idx in zip(src_order.tolist(), goal_order.tolist()):
+        sx0, sy0, sx1, sy1 = src_boxes[s_idx].tolist()
+        gx0, gy0, gx1, gy1 = goal_boxes[g_idx].tolist()
+
+        tile = src.crop((sx0, sy0, sx1, sy1))
+        goal_w = max(gx1 - gx0, 1)
+        goal_h = max(gy1 - gy0, 1)
+        if tile.size != (goal_w, goal_h):
+            tile = tile.resize((goal_w, goal_h), Image.Resampling.BILINEAR)
+        final.paste(tile, (gx0, gy0))
+    return final