DuOtto
136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
import cv2, numpy as np
|
||
from config import LOWER_SHAPE, UPPER_SHAPE
|
||
import cv2
|
||
import numpy as np
|
||
import time
|
||
import fitz # PyMuPDF for PDF rendering
|
||
|
||
|
||
|
||
def find_hud_placement(frame_vis):
|
||
# ----- Arrow Shape Detection & HUD Placement -----
|
||
# 1) Build a clean orange mask
|
||
hsv = cv2.cvtColor(frame_vis, cv2.COLOR_BGR2HSV)
|
||
mask = cv2.inRange(hsv, LOWER_SHAPE, UPPER_SHAPE)
|
||
kern = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))
|
||
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kern)
|
||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kern)
|
||
|
||
# 2) Find and filter contours
|
||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
for cnt in contours:
|
||
area = cv2.contourArea(cnt)
|
||
if area < 1000:
|
||
continue
|
||
|
||
# Approximate to polygon and require exactly 5 corners
|
||
peri = cv2.arcLength(cnt, True)
|
||
approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
|
||
if len(approx) != 5:
|
||
continue
|
||
|
||
pts = approx.reshape(-1,2)
|
||
centroid = np.mean(pts, axis=0)
|
||
|
||
# Find the arrow tip as the corner farthest from centroid
|
||
dists = [np.linalg.norm(pt - centroid) for pt in pts]
|
||
tip_pt = pts[int(np.argmax(dists))]
|
||
arrow_tip = (int(tip_pt[0]), int(tip_pt[1]))
|
||
|
||
# Place the HUD box in the arrow’s pointing direction
|
||
dir_vec = tip_pt - centroid
|
||
|
||
#ToDo remove
|
||
cv2.drawContours(frame_vis, [pts], -1, (0,255,255), 2)
|
||
|
||
return arrow_tip, dir_vec
|
||
|
||
# ——————————————————————————————————————————————
|
||
# PDF loading / rasterization at import time
|
||
# ——————————————————————————————————————————————
|
||
PDF_PATH = "DarkAngles.pdf" # change this to your file
|
||
PDF_DPI = 600 # controls resolution of rasterization
|
||
|
||
_doc = fitz.open(PDF_PATH)
|
||
_pdf_pages = []
|
||
for page in _doc:
|
||
# render page to pixmap at desired zoom
|
||
zoom = PDF_DPI / 72.0
|
||
mat = fitz.Matrix(zoom, zoom)
|
||
pix = page.get_pixmap(matrix=mat, alpha=False)
|
||
# convert pixmap to ndarray
|
||
img = np.frombuffer(pix.samples, dtype=np.uint8)
|
||
img = img.reshape(pix.height, pix.width, pix.n)
|
||
if pix.n == 4:
|
||
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
|
||
_pdf_pages.append(img)
|
||
_doc.close()
|
||
|
||
# ——————————————————————————————————————————————
|
||
# HUD drawing routines
|
||
# ——————————————————————————————————————————————
|
||
|
||
# rotated rectangle → 4 pts
|
||
def _rect_to_pts(center, size, angle_rad):
|
||
cx, cy = center
|
||
w, h = size
|
||
# local corners
|
||
pts = np.array([
|
||
[-w/2, -h/2],
|
||
[ w/2, -h/2],
|
||
[ w/2, h/2],
|
||
[-w/2, h/2],
|
||
])
|
||
# rotation
|
||
c, s = np.cos(angle_rad), np.sin(angle_rad)
|
||
R = np.array([[c, -s],[s, c]])
|
||
pts = pts.dot(R.T)
|
||
pts += np.array([cx, cy])
|
||
return pts.astype(np.float32)
|
||
|
||
def draw_hud_box(frame, tip, angle_rad, size=(120, 60), color=(255,0,255), thickness=2):
|
||
"""Draw a rotated HUD rectangle at `tip` pointing along `angle_rad`."""
|
||
# compute box corners
|
||
pts = _rect_to_pts(center=tip, size=size, angle_rad=angle_rad)
|
||
cv2.drawContours(frame, [pts.astype(int)], -1, color, thickness)
|
||
# label
|
||
cv2.putText(frame, "HUD", (int(tip[0]+5), int(tip[1]+5)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
|
||
return pts # return the quad for PDF warping
|
||
|
||
def draw_pdf_page(frame, page_index, dst_quad):
|
||
"""
|
||
Warp PDF page image #page_index into the quadrilateral dst_quad.
|
||
dst_quad: 4×2 float32 array of destination corners in clock-wise order.
|
||
"""
|
||
if page_index < 0 or page_index >= len(_pdf_pages):
|
||
return
|
||
src = _pdf_pages[page_index]
|
||
h, w = src.shape[:2]
|
||
|
||
# source corners (tl, tr, br, bl)
|
||
src_quad = np.array([[0,0], [w,0], [w,h], [0,h]], dtype=np.float32)
|
||
# compute homography
|
||
M = cv2.getPerspectiveTransform(src_quad, dst_quad)
|
||
# warp PDF page into place (transparent where outside)
|
||
warp = cv2.warpPerspective(src, M, (frame.shape[1], frame.shape[0]))
|
||
mask = cv2.warpPerspective(np.ones((h,w), dtype=np.uint8)*255, M, (frame.shape[1], frame.shape[0]))
|
||
# composite onto frame
|
||
inv = cv2.bitwise_not(mask)
|
||
bg = cv2.bitwise_and(frame, frame, mask=inv)
|
||
fg = cv2.bitwise_and(warp, warp, mask=mask)
|
||
np.copyto(frame, bg+fg)
|
||
|
||
# hud.py
|
||
class HUD:
|
||
def __init__(self, marker_id, pdf_pages, default_page=0):
|
||
self.id = marker_id
|
||
self.page_index = default_page
|
||
self.pdf_pages = pdf_pages # list of preloaded page images
|
||
self.ref_quad = None # the 4-corner quad where to draw the PDF
|
||
self.last_seen = 0 # for timeout if arrow goes away
|
||
|
||
def update(self, marker):
|
||
return
|
||
|
||
def draw(self, frame):
|
||
return |