Update app.py

app.py

@@ -1,70 +1,72 @@
+import gradio as gr
 import cv2
 import numpy as np
-import gradio as gr
 import heapq
 
+def compute_edge_costs(image):
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    gradient = cv2.Laplacian(gray, cv2.CV_64F)
+    edge_costs = np.abs(gradient)
+    return edge_costs
 
-def compute_cost_map(image):
-    grayscale = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
-    sobelx = cv2.Sobel(grayscale, cv2.CV_64F, 1, 0, ksize=5)
-    sobely = cv2.Sobel(grayscale, cv2.CV_64F, 0, 1, ksize=5)
-    magnitude = cv2.magnitude(sobelx, sobely)
-    return np.max(magnitude) - magnitude
+def dijkstra(edge_costs, start_point, end_point):
+    h, w = edge_costs.shape
+    visited = np.zeros((h, w), dtype=bool)
+    distances = np.full((h, w), float('inf'))
+    distances[start_point[1], start_point[0]] = 0
+    pq = [(0, start_point)]
+    
+    while pq:
+        curr_dist, (x, y) = heapq.heappop(pq)
 
-def neighbors(point, cost_map):
-    x, y = point
-    rows, cols = cost_map.shape
-    deltas = [(0, 1), (1, 0), (-1, 0), (0, -1), (1, 1), (-1, -1), (1, -1), (-1, 1)]
-    neigh = [(x + dx, y + dy) for dx, dy in deltas]
-    valid = [(x, y) for x, y in neigh if 0 <= x < rows and 0 <= y < cols]
-    return valid
+        if visited[y, x]:
+            continue
+        visited[y, x] = True
 
-def compute_shortest_path(cost_map, start, end):
-    rows, cols = cost_map.shape
-    distance = np.full((rows, cols), np.inf)
-    distance[start] = 0
-    prev = {}
-    priority_queue = [(0, start)]
-
-    while priority_queue:
-        current_distance, current_point = heapq.heappop(priority_queue)
-        if current_point == end:
-            break
-        for neighbor in neighbors(current_point, cost_map):
-            alt_distance = current_distance + cost_map[neighbor]
-            if alt_distance < distance[neighbor]:
-                distance[neighbor] = alt_distance
-                prev[neighbor] = current_point
-                heapq.heappush(priority_queue, (alt_distance, neighbor))
+        for dx in [-1, 0, 1]:
+            for dy in [-1, 0, 1]:
+                nx, ny = x + dx, y + dy
+                if 0 <= nx < w and 0 <= ny < h:
+                    new_dist = curr_dist + edge_costs[ny, nx]
+                    if new_dist < distances[ny, nx]:
+                        distances[ny, nx] = new_dist
+                        heapq.heappush(pq, (new_dist, (nx, ny)))
 
+    # Backtrack to find path from end to start
     path = []
-    while current_point in prev:
-        path.append(current_point)
-        current_point = prev[current_point]
-    path.append(start)
-    path.reverse()
-
+    curr_point = end_point
+    while curr_point != start_point:
+        x, y = curr_point
+        path.append(curr_point)
+        min_dist = distances[y, x]
+        for dx in [-1, 0, 1]:
+            for dy in [-1, 0, 1]:
+                nx, ny = x + dx, y + dy
+                if 0 <= nx < w and 0 <= ny < h and distances[ny, nx] < min_dist:
+                    min_dist = distances[ny, nx]
+                    curr_point = (nx, ny)
+    path.append(start_point)
+    
     return path
 
-def live_wire_segmentation(image, points):
-    cost_map = compute_cost_map(image)
-    path = []
-    for i in range(1, len(points)):
-        segment = compute_shortest_path(cost_map, points[i-1], points[i])
-        path.extend(segment)
-
-    # Draw the path on the image
-    for i in range(1, len(path)):
-        cv2.line(image, tuple(path[i-1]), tuple(path[i]), (0, 255, 0), 2)
-    return image
+def segment_image(image, points):
+    edge_costs = compute_edge_costs(image)
+    
+    start_point = tuple(map(int, points[0]))
+    end_point = tuple(map(int, points[-1]))
+    
+    path = dijkstra(edge_costs, start_point, end_point)
+    
+    mask = np.zeros_like(image)
+    for x, y in path:
+        cv2.circle(mask, (x, y), 2, (0, 255, 0), -1)
+    return mask
 
-iface = gr.Interface(
-    fn=live_wire_segmentation, 
-    inputs=[gr.inputs.Image(shape=(512, 512), label="Original Image"), 
-            gr.inputs.Sketchpad()],
-    outputs=gr.outputs.Image()
+interface = gr.Interface(
+    segment_image,
+    inputs=[gr.inputs.Image(shape=(200, 200)), gr.inputs.Sketchpad()],
+    outputs=gr.outputs.Image(shape=(200, 200)),
+    live=True
 )
 
-iface.launch()
-
-
+interface.launch()