phase 4
This commit is contained in:
1
core/gpu/models/cv/__init__.py
Normal file
1
core/gpu/models/cv/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""CV operations — pure OpenCV, no ML models."""
|
||||
258
core/gpu/models/cv/edges.py
Normal file
258
core/gpu/models/cv/edges.py
Normal file
@@ -0,0 +1,258 @@
|
||||
"""
|
||||
Edge detection — Canny + HoughLinesP → parallel line pairs → bounding boxes.
|
||||
|
||||
Finds horizontal line pairs with consistent spacing, which correspond to
|
||||
the top and bottom edges of advertising hoardings.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def detect_edges(
|
||||
image: np.ndarray,
|
||||
canny_low: int = 50,
|
||||
canny_high: int = 150,
|
||||
hough_threshold: int = 80,
|
||||
hough_min_length: int = 100,
|
||||
hough_max_gap: int = 10,
|
||||
pair_max_distance: int = 200,
|
||||
pair_min_distance: int = 15,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Find horizontal line pairs that likely bound advertising hoardings.
|
||||
|
||||
Returns list of dicts with keys: x, y, w, h, confidence, label.
|
||||
Each box represents the region between a detected pair of parallel
|
||||
horizontal lines.
|
||||
"""
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
||||
edges = cv2.Canny(gray, canny_low, canny_high)
|
||||
|
||||
raw_lines = cv2.HoughLinesP(
|
||||
edges,
|
||||
rho=1,
|
||||
theta=np.pi / 180,
|
||||
threshold=hough_threshold,
|
||||
minLineLength=hough_min_length,
|
||||
maxLineGap=hough_max_gap,
|
||||
)
|
||||
|
||||
if raw_lines is None:
|
||||
return []
|
||||
|
||||
# Filter to near-horizontal lines (within 10 degrees)
|
||||
horizontals = _filter_horizontal(raw_lines, max_angle_deg=10)
|
||||
|
||||
if len(horizontals) < 2:
|
||||
return []
|
||||
|
||||
# Find pairs of parallel horizontals with consistent spacing
|
||||
pairs = _find_line_pairs(
|
||||
horizontals,
|
||||
min_distance=pair_min_distance,
|
||||
max_distance=pair_max_distance,
|
||||
)
|
||||
|
||||
# Convert pairs to bounding boxes
|
||||
h, w = image.shape[:2]
|
||||
results = []
|
||||
for top_line, bottom_line in pairs:
|
||||
box = _pair_to_bbox(top_line, bottom_line, frame_width=w, frame_height=h)
|
||||
if box is not None:
|
||||
results.append(box)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _filter_horizontal(lines: np.ndarray, max_angle_deg: float = 10) -> list[tuple]:
|
||||
"""Keep only lines within max_angle_deg of horizontal."""
|
||||
max_slope = np.tan(np.radians(max_angle_deg))
|
||||
result = []
|
||||
for line in lines:
|
||||
x1, y1, x2, y2 = line[0]
|
||||
dx = x2 - x1
|
||||
if dx == 0:
|
||||
continue
|
||||
slope = abs((y2 - y1) / dx)
|
||||
if slope <= max_slope:
|
||||
y_mid = (y1 + y2) / 2
|
||||
x_min = min(x1, x2)
|
||||
x_max = max(x1, x2)
|
||||
length = np.sqrt(dx**2 + (y2 - y1) ** 2)
|
||||
result.append((x_min, x_max, y_mid, length))
|
||||
return result
|
||||
|
||||
|
||||
def _find_line_pairs(
|
||||
horizontals: list[tuple],
|
||||
min_distance: int,
|
||||
max_distance: int,
|
||||
) -> list[tuple]:
|
||||
"""
|
||||
Find pairs of horizontal lines that could be top/bottom of a hoarding.
|
||||
|
||||
Lines must overlap horizontally and be spaced within [min_distance, max_distance].
|
||||
"""
|
||||
# Sort by y position
|
||||
sorted_lines = sorted(horizontals, key=lambda l: l[2])
|
||||
|
||||
pairs = []
|
||||
used = set()
|
||||
|
||||
for i, top in enumerate(sorted_lines):
|
||||
if i in used:
|
||||
continue
|
||||
for j, bottom in enumerate(sorted_lines[i + 1 :], start=i + 1):
|
||||
if j in used:
|
||||
continue
|
||||
|
||||
y_gap = bottom[2] - top[2]
|
||||
if y_gap < min_distance:
|
||||
continue
|
||||
if y_gap > max_distance:
|
||||
break # sorted by y, no point checking further
|
||||
|
||||
# Check horizontal overlap
|
||||
overlap_start = max(top[0], bottom[0])
|
||||
overlap_end = min(top[1], bottom[1])
|
||||
overlap = overlap_end - overlap_start
|
||||
|
||||
# Require at least 50% overlap relative to shorter line
|
||||
shorter_length = min(top[1] - top[0], bottom[1] - bottom[0])
|
||||
if shorter_length > 0 and overlap / shorter_length >= 0.5:
|
||||
pairs.append((top, bottom))
|
||||
used.add(i)
|
||||
used.add(j)
|
||||
break
|
||||
|
||||
return pairs
|
||||
|
||||
|
||||
def _pair_to_bbox(
|
||||
top: tuple,
|
||||
bottom: tuple,
|
||||
frame_width: int,
|
||||
frame_height: int,
|
||||
) -> dict | None:
|
||||
"""Convert a line pair to a bounding box dict."""
|
||||
x = int(max(0, min(top[0], bottom[0])))
|
||||
y = int(max(0, top[2]))
|
||||
x2 = int(min(frame_width, max(top[1], bottom[1])))
|
||||
y2 = int(min(frame_height, bottom[2]))
|
||||
w = x2 - x
|
||||
h = y2 - y
|
||||
|
||||
if w < 20 or h < 5:
|
||||
return None
|
||||
|
||||
# Confidence based on line lengths relative to box width
|
||||
avg_line_length = (top[3] + bottom[3]) / 2
|
||||
coverage = min(1.0, avg_line_length / max(w, 1))
|
||||
|
||||
return {
|
||||
"x": x,
|
||||
"y": y,
|
||||
"w": w,
|
||||
"h": h,
|
||||
"confidence": round(coverage, 3),
|
||||
"label": "edge_region",
|
||||
}
|
||||
|
||||
|
||||
def _np_to_b64_jpeg(image: np.ndarray, quality: int = 70) -> str:
|
||||
"""Encode a numpy image (BGR or grayscale) as base64 JPEG."""
|
||||
ok, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
||||
if not ok:
|
||||
return ""
|
||||
return base64.b64encode(buf.tobytes()).decode()
|
||||
|
||||
|
||||
def detect_edges_debug(
|
||||
image: np.ndarray,
|
||||
canny_low: int = 50,
|
||||
canny_high: int = 150,
|
||||
hough_threshold: int = 80,
|
||||
hough_min_length: int = 100,
|
||||
hough_max_gap: int = 10,
|
||||
pair_max_distance: int = 200,
|
||||
pair_min_distance: int = 15,
|
||||
) -> dict:
|
||||
"""
|
||||
Same as detect_edges but returns intermediate visualizations.
|
||||
|
||||
Returns dict with:
|
||||
regions: list[dict] — same boxes as detect_edges
|
||||
edge_overlay_b64: str — Canny edge image as base64 JPEG
|
||||
lines_overlay_b64: str — frame with Hough lines drawn
|
||||
horizontal_count: int — number of horizontal lines found
|
||||
pair_count: int — number of line pairs found
|
||||
"""
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
|
||||
edges = cv2.Canny(gray, canny_low, canny_high)
|
||||
|
||||
# Edge overlay — Canny output as-is (white edges on black)
|
||||
edge_overlay_b64 = _np_to_b64_jpeg(edges)
|
||||
|
||||
raw_lines = cv2.HoughLinesP(
|
||||
edges,
|
||||
rho=1,
|
||||
theta=np.pi / 180,
|
||||
threshold=hough_threshold,
|
||||
minLineLength=hough_min_length,
|
||||
maxLineGap=hough_max_gap,
|
||||
)
|
||||
|
||||
# Lines overlay — draw all Hough lines on a copy of the frame
|
||||
lines_vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
if raw_lines is not None:
|
||||
for line in raw_lines:
|
||||
x1, y1, x2, y2 = line[0]
|
||||
cv2.line(lines_vis, (x1, y1), (x2, y2), (0, 0, 255), 1)
|
||||
|
||||
horizontals = []
|
||||
if raw_lines is not None:
|
||||
horizontals = _filter_horizontal(raw_lines, max_angle_deg=10)
|
||||
|
||||
# Draw horizontal lines in cyan, thicker
|
||||
for h_line in horizontals:
|
||||
x_min, x_max, y_mid, _ = h_line
|
||||
cv2.line(lines_vis, (int(x_min), int(y_mid)), (int(x_max), int(y_mid)), (255, 255, 0), 2)
|
||||
|
||||
pairs = []
|
||||
if len(horizontals) >= 2:
|
||||
pairs = _find_line_pairs(
|
||||
horizontals,
|
||||
min_distance=pair_min_distance,
|
||||
max_distance=pair_max_distance,
|
||||
)
|
||||
|
||||
# Draw paired lines in green
|
||||
for top_line, bottom_line in pairs:
|
||||
cv2.line(lines_vis, (int(top_line[0]), int(top_line[2])),
|
||||
(int(top_line[1]), int(top_line[2])), (0, 255, 0), 2)
|
||||
cv2.line(lines_vis, (int(bottom_line[0]), int(bottom_line[2])),
|
||||
(int(bottom_line[1]), int(bottom_line[2])), (0, 255, 0), 2)
|
||||
|
||||
lines_overlay_b64 = _np_to_b64_jpeg(lines_vis)
|
||||
|
||||
# Build region boxes (same logic as detect_edges)
|
||||
h, w = image.shape[:2]
|
||||
regions = []
|
||||
for top_line, bottom_line in pairs:
|
||||
box = _pair_to_bbox(top_line, bottom_line, frame_width=w, frame_height=h)
|
||||
if box is not None:
|
||||
regions.append(box)
|
||||
|
||||
return {
|
||||
"regions": regions,
|
||||
"edge_overlay_b64": edge_overlay_b64,
|
||||
"lines_overlay_b64": lines_overlay_b64,
|
||||
"horizontal_count": len(horizontals),
|
||||
"pair_count": len(pairs),
|
||||
}
|
||||
86
core/gpu/models/cv/segmentation.py
Normal file
86
core/gpu/models/cv/segmentation.py
Normal file
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
Field segmentation — HSV green mask → pitch boundary contour.
|
||||
|
||||
Pure OpenCV. Called by the inference server endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def segment_field(
|
||||
image: np.ndarray,
|
||||
hue_low: int = 30,
|
||||
hue_high: int = 85,
|
||||
sat_low: int = 30,
|
||||
sat_high: int = 255,
|
||||
val_low: int = 30,
|
||||
val_high: int = 255,
|
||||
morph_kernel: int = 15,
|
||||
min_area_ratio: float = 0.05,
|
||||
) -> dict:
|
||||
"""
|
||||
Detect the pitch area using HSV green thresholding.
|
||||
|
||||
Returns dict with:
|
||||
boundary: list of [x, y] points
|
||||
coverage: float (fraction of frame)
|
||||
"""
|
||||
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
|
||||
|
||||
lower = np.array([hue_low, sat_low, val_low])
|
||||
upper = np.array([hue_high, sat_high, val_high])
|
||||
mask = cv2.inRange(hsv, lower, upper)
|
||||
|
||||
k = morph_kernel
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
|
||||
|
||||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
h, w = image.shape[:2]
|
||||
min_area = min_area_ratio * h * w
|
||||
boundary = []
|
||||
coverage = 0.0
|
||||
|
||||
if contours:
|
||||
large = [c for c in contours if cv2.contourArea(c) >= min_area]
|
||||
if large:
|
||||
pitch_contour = max(large, key=cv2.contourArea)
|
||||
boundary = pitch_contour.reshape(-1, 2).tolist()
|
||||
coverage = cv2.contourArea(pitch_contour) / (h * w)
|
||||
|
||||
refined = np.zeros_like(mask)
|
||||
cv2.drawContours(refined, [pitch_contour], -1, 255, cv2.FILLED)
|
||||
mask = refined
|
||||
|
||||
return {
|
||||
"boundary": boundary,
|
||||
"coverage": coverage,
|
||||
"mask": mask,
|
||||
}
|
||||
|
||||
|
||||
def segment_field_debug(
|
||||
image: np.ndarray,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
"""Same as segment_field but includes a mask overlay for the editor."""
|
||||
result = segment_field(image, **kwargs)
|
||||
mask = result["mask"]
|
||||
|
||||
# RGBA overlay: solid green where mask, fully transparent elsewhere
|
||||
h, w = image.shape[:2]
|
||||
overlay = np.zeros((h, w, 4), dtype=np.uint8)
|
||||
overlay[mask > 0] = [0, 255, 0, 255]
|
||||
_, buf = cv2.imencode(".png", overlay)
|
||||
result["mask_overlay_b64"] = base64.b64encode(buf.tobytes()).decode()
|
||||
|
||||
# Don't send the raw mask over HTTP
|
||||
del result["mask"]
|
||||
return result
|
||||
Reference in New Issue
Block a user