Spaces:
Sleeping
Sleeping
File size: 6,937 Bytes
8bbb872 | 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | from __future__ import annotations
from pathlib import Path
import os
import cv2
import numpy as np
from ultralytics import YOLO
def list_images(folder: Path):
exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
return sorted([p for p in folder.iterdir() if p.suffix.lower() in exts])
def find_weights(project_root: Path) -> Path | None:
candidates = [
project_root / "weights" / "best.pt",
project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "best.pt",
project_root / "runs" / "classify" / "runs_cls" / "eye_open_closed_cpu" / "weights" / "last.pt",
project_root / "runs_cls" / "eye_open_closed_cpu" / "weights" / "best.pt",
project_root / "runs_cls" / "eye_open_closed_cpu" / "weights" / "last.pt",
]
return next((p for p in candidates if p.is_file()), None)
def detect_eyelid_boundary(gray: np.ndarray) -> np.ndarray | None:
"""
Returns an ellipse fit to the largest contour near the eye boundary.
Output format: (center(x,y), (axis1, axis2), angle) or None.
"""
blur = cv2.GaussianBlur(gray, (5, 5), 0)
edges = cv2.Canny(blur, 40, 120)
edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1)
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return None
contours = sorted(contours, key=cv2.contourArea, reverse=True)
for c in contours:
if len(c) >= 5 and cv2.contourArea(c) > 50:
return cv2.fitEllipse(c)
return None
def detect_pupil_center(gray: np.ndarray) -> tuple[int, int] | None:
"""
More robust pupil detection:
- enhance contrast (CLAHE)
- find dark blobs
- score by circularity and proximity to center
"""
h, w = gray.shape
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
eq = clahe.apply(gray)
blur = cv2.GaussianBlur(eq, (7, 7), 0)
# Focus on the central region to avoid eyelashes/edges
cx, cy = w // 2, h // 2
rx, ry = int(w * 0.3), int(h * 0.3)
x0, x1 = max(cx - rx, 0), min(cx + rx, w)
y0, y1 = max(cy - ry, 0), min(cy + ry, h)
roi = blur[y0:y1, x0:x1]
# Inverted threshold to capture dark pupil
_, thresh = cv2.threshold(roi, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8), iterations=2)
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8), iterations=1)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return None
best = None
best_score = -1.0
for c in contours:
area = cv2.contourArea(c)
if area < 15:
continue
perimeter = cv2.arcLength(c, True)
if perimeter <= 0:
continue
circularity = 4 * np.pi * (area / (perimeter * perimeter))
if circularity < 0.3:
continue
m = cv2.moments(c)
if m["m00"] == 0:
continue
px = int(m["m10"] / m["m00"]) + x0
py = int(m["m01"] / m["m00"]) + y0
# Score by circularity and distance to center
dist = np.hypot(px - cx, py - cy) / max(w, h)
score = circularity - dist
if score > best_score:
best_score = score
best = (px, py)
return best
def is_focused(pupil_center: tuple[int, int], img_shape: tuple[int, int]) -> bool:
"""
Decide focus based on pupil offset from image center.
"""
h, w = img_shape
cx, cy = w // 2, h // 2
px, py = pupil_center
dx = abs(px - cx) / max(w, 1)
dy = abs(py - cy) / max(h, 1)
return (dx < 0.12) and (dy < 0.12)
def annotate(img_bgr: np.ndarray, ellipse, pupil_center, focused: bool, cls_label: str, conf: float):
out = img_bgr.copy()
if ellipse is not None:
cv2.ellipse(out, ellipse, (0, 255, 255), 2)
if pupil_center is not None:
cv2.circle(out, pupil_center, 4, (0, 0, 255), -1)
label = f"{cls_label} ({conf:.2f}) | focused={int(focused)}"
cv2.putText(out, label, (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
return out
def main():
project_root = Path(__file__).resolve().parent.parent
data_dir = project_root / "Dataset"
alt_data_dir = project_root / "DATA"
out_dir = project_root / "runs_focus"
out_dir.mkdir(parents=True, exist_ok=True)
weights = find_weights(project_root)
if weights is None:
print("Weights not found. Train first.")
return
# Support both Dataset/test/{open,closed} and Dataset/{open,closed}
def resolve_test_dirs(root: Path):
test_open = root / "test" / "open"
test_closed = root / "test" / "closed"
if test_open.exists() and test_closed.exists():
return test_open, test_closed
test_open = root / "open"
test_closed = root / "closed"
if test_open.exists() and test_closed.exists():
return test_open, test_closed
alt_closed = root / "close"
if test_open.exists() and alt_closed.exists():
return test_open, alt_closed
return None, None
test_open, test_closed = resolve_test_dirs(data_dir)
if (test_open is None or test_closed is None) and alt_data_dir.exists():
test_open, test_closed = resolve_test_dirs(alt_data_dir)
if not test_open.exists() or not test_closed.exists():
print("Test folders missing. Expected:")
print(test_open)
print(test_closed)
return
test_files = list_images(test_open) + list_images(test_closed)
print("Total test images:", len(test_files))
max_images = int(os.getenv("MAX_IMAGES", "0"))
if max_images > 0:
test_files = test_files[:max_images]
print("Limiting to MAX_IMAGES:", max_images)
model = YOLO(str(weights))
results = model.predict(test_files, imgsz=224, device="cpu", verbose=False)
names = model.names
for r in results:
probs = r.probs
top_idx = int(probs.top1)
top_conf = float(probs.top1conf)
pred_label = names[top_idx]
img = cv2.imread(r.path)
if img is None:
continue
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
ellipse = detect_eyelid_boundary(gray)
pupil_center = detect_pupil_center(gray)
focused = False
if pred_label.lower() == "open" and pupil_center is not None:
focused = is_focused(pupil_center, gray.shape)
annotated = annotate(img, ellipse, pupil_center, focused, pred_label, top_conf)
out_path = out_dir / (Path(r.path).stem + "_annotated.jpg")
cv2.imwrite(str(out_path), annotated)
print(f"{Path(r.path).name}: pred={pred_label} conf={top_conf:.3f} focused={focused}")
print(f"\nAnnotated outputs saved to: {out_dir}")
if __name__ == "__main__":
main()
|