86 lines
3.1 KiB
Python
86 lines
3.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
def bury_pipe_stream():
|
|
src_path = '/Users/davidkotnik/.gemini/antigravity/brain/07019d04-a214-43ab-9565-86f4e8f17e5b/uploaded_media_1769607894587.jpg'
|
|
|
|
print(f"Loading {src_path}")
|
|
img = cv2.imread(src_path)
|
|
if img is None: return
|
|
|
|
# 1. GrabCut to isolate object from BG
|
|
mask = np.zeros(img.shape[:2], np.uint8)
|
|
bgdModel = np.zeros((1,65),np.float64)
|
|
fgdModel = np.zeros((1,65),np.float64)
|
|
h, w = img.shape[:2]
|
|
cv2.grabCut(img, mask, (10, 10, w-20, h-20), bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
|
|
mask_final = np.where((mask==2)|(mask==0),0,1).astype('uint8')
|
|
|
|
# 2. REMOVE OUTER WALLS (The "Burying" Step)
|
|
# The walls are the darkest parts of the ground structure.
|
|
# The "Top Rim" (ground level) is lighter brown.
|
|
# The water is murky brown/gray.
|
|
|
|
# Convert to HSV to separate by brightness/color
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
|
|
# Define "Wall" Color range (Dark Brown / Black shadow)
|
|
# H: Any (Browns are 0-20 or 160-180)
|
|
# S: Low to Medium
|
|
# V: VERY LOW (Shadows)
|
|
|
|
# Let's visualize: brightness < 40?
|
|
v_channel = hsv[:,:,2]
|
|
|
|
# Create a mask of "Very Dark" pixels
|
|
wall_mask = (v_channel < 60)
|
|
|
|
# But we want to keep outlines (which are thin dark lines).
|
|
# Walls are LARGE blocks of dark.
|
|
# We can use morphological opening to select only LARGE dark areas (walls) and ignore thin lines (outlines).
|
|
wall_mask_uint8 = wall_mask.astype(np.uint8)
|
|
kernel = np.ones((5,5), np.uint8)
|
|
# Erode then Dilate -> Removes small noise (lines), keeps big blobs (walls)
|
|
walls_only = cv2.morphologyEx(wall_mask_uint8, cv2.MORPH_OPEN, kernel, iterations=2)
|
|
|
|
# Now, subtract 'walls_only' from our main mask.
|
|
# We want mask_final to be 0 where walls_only is 1.
|
|
mask_buried = mask_final.copy()
|
|
mask_buried[walls_only == 1] = 0
|
|
|
|
# 3. FIX: This might delete the interior of the pipe (shadow) or the drain hole!
|
|
# We must PROTECT the pipe/drain area.
|
|
# The Drain is at the top right. Let's define a Region of Interest (ROI) that we DO NOT touch.
|
|
# Pipe is roughly in top right quadrant.
|
|
# Heuristic: Only apply wall removal to the bottom half/rim?
|
|
# Or better: The walls are on the OUTSIDE boundary.
|
|
|
|
# Let's perform the subtraction.
|
|
|
|
# 4. FIX PIPE HOLE (Make opaque again if GrabCut lost it)
|
|
# (Re-using logic from previous step if needed, but GrabCut usually keeps it).
|
|
|
|
# Combine
|
|
b, g, r = cv2.split(img)
|
|
alpha = mask_buried * 255
|
|
|
|
img_rgba = cv2.merge([b, g, r, alpha])
|
|
|
|
# Crop
|
|
coords = cv2.findNonZero(alpha)
|
|
if coords is not None:
|
|
x, y, cw, ch = cv2.boundingRect(coords)
|
|
img_rgba = img_rgba[y:y+ch, x:x+cw]
|
|
|
|
targets = [
|
|
'/Users/davidkotnik/repos/novafarma/main/assets/stream_pipe.png',
|
|
'/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Environment/stream_pipe.png'
|
|
]
|
|
for t in targets:
|
|
cv2.imwrite(t, img_rgba)
|
|
print(f"Saved {t}")
|
|
|
|
if __name__ == "__main__":
|
|
bury_pipe_stream()
|