Spaces:
Build error
Build error
File size: 3,837 Bytes
4bb934b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import numpy as np
import cv2
import heapq
import matplotlib.pyplot as plt
from collections import deque
# 1. Compute local minima as markers
def get_local_minima(gray):
kernel = np.ones((3, 3), np.uint8)
eroded = cv2.erode(gray, kernel)
minima = (gray == eroded)
return minima.astype(np.uint8)
# 2. Label each connected component (marker)
def label_markers(minima):
num_labels, markers = cv2.connectedComponents(minima)
return markers, num_labels
# 3. Watershed from scratch
def watershed_from_scratch(gray, markers):
h, w = gray.shape
# Constants
WATERSHED = -1
INIT = -2
# Initialize label and visited map
label_map = np.full((h, w), INIT, dtype=np.int32)
label_map[markers > 0] = markers[markers > 0]
# Priority queue for pixels: (intensity, y, x)
pq = []
# Populate queue with boundary of initial markers
for y in range(h):
for x in range(w):
if markers[y, x] > 0:
for dy in [-1, 0, 1]:
for dx in [-1, 0, 1]:
ny, nx = y + dy, x + dx
if 0 <= ny < h and 0 <= nx < w:
if markers[ny, nx] == 0 and label_map[ny, nx] == INIT:
heapq.heappush(pq, (gray[ny, nx], ny, nx))
label_map[ny, nx] = 0 # Mark as in queue
# Flooding
while pq:
intensity, y, x = heapq.heappop(pq)
neighbor_labels = set()
for dy in [-1, 0, 1]:
for dx in [-1, 0, 1]:
ny, nx = y + dy, x + dx
if 0 <= ny < h and 0 <= nx < w:
lbl = label_map[ny, nx]
if lbl > 0:
neighbor_labels.add(lbl)
if len(neighbor_labels) == 1:
label_map[y, x] = neighbor_labels.pop()
elif len(neighbor_labels) > 1:
label_map[y, x] = WATERSHED
# Add unvisited neighbors to the queue
for dy in [-1, 0, 1]:
for dx in [-1, 0, 1]:
ny, nx = y + dy, x + dx
if 0 <= ny < h and 0 <= nx < w:
if label_map[ny, nx] == INIT:
heapq.heappush(pq, (gray[ny, nx], ny, nx))
label_map[ny, nx] = 0 # Mark as in queue
return label_map
def generate_watershed(iamge_path):
# Load grayscale image
image = cv2.imread(iamge_path, cv2.IMREAD_GRAYSCALE)
image = cv2.GaussianBlur(image, (5, 5), 0)
minima = get_local_minima(image)
markers, num_labels = label_markers(minima)
result = watershed_from_scratch(image, markers)
# Visualization
output = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)
output[result == -1] = [255, 0, 0] # Watershed lines in red
output[result > 0] = [0, 255, 0] # Segments in green
output[markers > 0] = [0, 0, 255] # Original minima in blue
return image,output
if __name__ == "__main__":
# Run the process
# Load grayscale image
image = cv2.imread("/home/akshat/projects/CSL7360_Project/bird.jpeg", cv2.IMREAD_GRAYSCALE)
image = cv2.GaussianBlur(image, (5, 5), 0)
minima = get_local_minima(image)
markers, num_labels = label_markers(minima)
result = watershed_from_scratch(image, markers)
# Visualization
output = np.zeros((image.shape[0], image.shape[1], 3), dtype=np.uint8)
output[result == -1] = [255, 0, 0] # Watershed lines in red
output[result > 0] = [0, 255, 0] # Segments in green
output[markers > 0] = [0, 0, 255] # Original minima in blue
# Save the original grayscale and the output image
cv2.imwrite("original_grayscale.png", image)
cv2.imwrite("watershed_output.png", output)
print("Images saved as 'original_grayscale.png' and 'watershed_output.png'")
|