| import gradio as gr |
| import cv2 |
| import numpy as np |
| from pathlib import Path |
| from ultralytics import YOLO |
|
|
| |
| MODEL_PATH = "yolov8_cattle_keypoints.pt" |
|
|
| KP_NAMES = [ |
| "left_ear_tip", "right_ear_tip", "left_ear_base", "right_ear_base", |
| "left_eye", "right_eye", "nose_left", "nose_right", "nose_tip", |
| "mouth_left", "mouth_right", "chin_left", "chin_right", |
| ] |
|
|
| KP_COLORS = [ |
| (255, 69, 0), |
| ( 30, 144, 255), |
| (255, 165, 0), |
| ( 0, 191, 255), |
| (154, 205, 50), |
| (238, 130, 238), |
| ( 0, 255, 127), |
| (255, 215, 0), |
| (255, 255, 0), |
| (255, 20, 147), |
| ( 0, 255, 255), |
| (255, 140, 0), |
| (147, 112, 219), |
| ] |
|
|
| SKELETON = [ |
| (0, 2), (1, 3), |
| (2, 4), (3, 5), |
| (4, 5), |
| (6, 8), (7, 8), |
| (6, 9), (7, 10), |
| (9, 10), |
| (9, 11), (10, 12), |
| (11, 12), |
| (4, 6), (5, 7), |
| ] |
|
|
| |
| _model = None |
|
|
| def load_model(): |
| global _model |
| if _model is None: |
| if not Path(MODEL_PATH).exists(): |
| raise FileNotFoundError( |
| f"Model file '{MODEL_PATH}' not found. " |
| "Upload yolov8_cattle_keypoints.pt to the Space root." |
| ) |
| _model = YOLO(MODEL_PATH) |
| return _model |
|
|
|
|
| |
| def run_inference(image: np.ndarray, conf_threshold: float, show_labels: bool): |
| if image is None: |
| return None, "Upload an image first." |
|
|
| m = load_model() |
| results = m.predict(source=image, conf=float(conf_threshold), verbose=False) |
|
|
| annotated = image.copy() |
| table_rows = [] |
|
|
| for result in results: |
| boxes = result.boxes |
| keypoints_data = result.keypoints |
| if boxes is None or keypoints_data is None or len(boxes) == 0: |
| continue |
|
|
| for det_idx in range(len(boxes)): |
| conf = float(boxes.conf[det_idx]) |
| kps = keypoints_data.data[det_idx].cpu().numpy() |
|
|
| |
| x1, y1, x2, y2 = boxes.xyxy[det_idx].cpu().numpy().astype(int) |
| cv2.rectangle(annotated, (x1, y1), (x2, y2), (255, 255, 255), 2) |
| cv2.putText( |
| annotated, f"cattle {conf:.2f}", |
| (x1, max(y1 - 8, 0)), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA, |
| ) |
|
|
| |
| for (i, j) in SKELETON: |
| if i >= len(kps) or j >= len(kps): |
| continue |
| xi, yi, vi = kps[i] |
| xj, yj, vj = kps[j] |
| if vi < 0.5 or vj < 0.5: |
| continue |
| cv2.line( |
| annotated, |
| (int(xi), int(yi)), (int(xj), int(yj)), |
| (180, 180, 180), 1, cv2.LINE_AA, |
| ) |
|
|
| |
| row = {"detection": det_idx + 1, "confidence": f"{conf:.3f}"} |
| for kp_idx, (kx, ky, kv) in enumerate(kps): |
| name = KP_NAMES[kp_idx] |
| color = KP_COLORS[kp_idx] |
| if kv > 0.5: |
| cv2.circle(annotated, (int(kx), int(ky)), 6, color, -1) |
| cv2.circle(annotated, (int(kx), int(ky)), 7, (0, 0, 0), 1) |
| if show_labels: |
| cv2.putText( |
| annotated, name, |
| (int(kx) + 8, int(ky) - 4), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.38, |
| color, 1, cv2.LINE_AA, |
| ) |
| row[name] = f"({int(kx)}, {int(ky)}) vis={kv:.2f}" |
| else: |
| row[name] = "not visible" |
|
|
| table_rows.append(row) |
|
|
| if table_rows: |
| lines = [f"### {len(table_rows)} detection(s) found\n"] |
| for row in table_rows: |
| lines.append(f"**Detection {row['detection']}** — conf {row['confidence']}") |
| for name in KP_NAMES: |
| lines.append(f" - `{name}`: {row.get(name, 'n/a')}") |
| lines.append("") |
| results_md = "\n".join(lines) |
| else: |
| results_md = "### No cattle detected\nTry lowering the confidence threshold." |
|
|
| return annotated, results_md |
|
|
|
|
| |
| CSS = """ |
| @import url('https://fonts.googleapis.com/css2?family=Bebas+Neue&family=DM+Mono:wght@400;500&family=DM+Sans:wght@300;400;500&display=swap'); |
| |
| :root { |
| --bg: #0d0f0e; |
| --surface: #161a18; |
| --border: #2a332e; |
| --accent: #4ffe9a; |
| --muted: #7a8c80; |
| --text: #e8ede9; |
| --radius: 4px; |
| } |
| body, .gradio-container { |
| background: var(--bg) !important; |
| font-family: 'DM Sans', sans-serif !important; |
| color: var(--text) !important; |
| } |
| .hdr { |
| padding: 2.5rem 0 1.5rem; |
| text-align: center; |
| border-bottom: 1px solid var(--border); |
| margin-bottom: 2rem; |
| } |
| .hdr h1 { |
| font-family: 'Bebas Neue', sans-serif; |
| font-size: clamp(2.8rem, 7vw, 5.5rem); |
| letter-spacing: 0.08em; |
| color: var(--accent); |
| margin: 0; |
| line-height: 1; |
| text-shadow: 0 0 40px rgba(79,254,154,0.25); |
| } |
| .hdr p { |
| font-family: 'DM Mono', monospace; |
| font-size: 0.78rem; |
| color: var(--muted); |
| letter-spacing: 0.15em; |
| text-transform: uppercase; |
| margin: 0.6rem 0 0; |
| } |
| .tags { |
| display: flex; gap: 0.5rem; flex-wrap: wrap; |
| justify-content: center; margin-top: 0.8rem; |
| } |
| .tag { |
| font-family: 'DM Mono', monospace; |
| font-size: 0.68rem; letter-spacing: 0.1em; |
| text-transform: uppercase; padding: 0.25rem 0.65rem; |
| border: 1px solid var(--border); border-radius: 2px; color: var(--muted); |
| } |
| .tag.hot { border-color: var(--accent); color: var(--accent); } |
| button.primary { |
| background: var(--accent) !important; |
| color: #0d0f0e !important; |
| font-family: 'Bebas Neue', sans-serif !important; |
| font-size: 1.1rem !important; |
| letter-spacing: 0.12em !important; |
| border: none !important; |
| border-radius: var(--radius) !important; |
| padding: 0.7rem 2rem !important; |
| transition: opacity 0.15s, transform 0.1s !important; |
| } |
| button.primary:hover { opacity: 0.85 !important; transform: translateY(-1px) !important; } |
| input[type=range] { accent-color: var(--accent) !important; } |
| input[type=checkbox] { accent-color: var(--accent) !important; } |
| """ |
|
|
| HEADER_HTML = """ |
| <div class="hdr"> |
| <h1>CattleFace · Pose</h1> |
| <p>YOLOv8 · 13-point facial landmark detection for bovines</p> |
| <div class="tags"> |
| <span class="tag hot">13 keypoints</span> |
| <span class="tag">ears · eyes · nose · mouth · chin</span> |
| <span class="tag hot">real-time inference</span> |
| <span class="tag">UARK-AICV benchmark</span> |
| </div> |
| </div> |
| """ |
|
|
| FOOTER_HTML = """ |
| <div style="text-align:center;padding:1.5rem 0 0.5rem; |
| font-family:'DM Mono',monospace;font-size:0.7rem; |
| color:#7a8c80;letter-spacing:0.08em;"> |
| MODEL · YOLOv8s-pose | |
| DATASET · UARK-AICV/CattleFace-RGBT-benchmark | |
| 13 FACIAL LANDMARKS |
| </div> |
| """ |
|
|
| |
| with gr.Blocks(title="CattleFace Pose") as demo: |
| gr.HTML(HEADER_HTML) |
|
|
| with gr.Row(): |
| with gr.Column(scale=1): |
| inp_image = gr.Image( |
| label="Input Image", |
| type="numpy", |
| sources=["upload", "webcam", "clipboard"], |
| ) |
| conf_slider = gr.Slider( |
| minimum=0.05, maximum=0.95, value=0.25, step=0.05, |
| label="Confidence Threshold", |
| ) |
| show_labels = gr.Checkbox(value=True, label="Show keypoint labels") |
| run_btn = gr.Button("Detect Landmarks", variant="primary") |
|
|
| with gr.Column(scale=1): |
| out_image = gr.Image(label="Annotated Output", type="numpy") |
| out_text = gr.Markdown(label="Keypoint Details") |
|
|
| run_btn.click( |
| fn=run_inference, |
| inputs=[inp_image, conf_slider, show_labels], |
| outputs=[out_image, out_text], |
| ) |
|
|
| gr.HTML(FOOTER_HTML) |
|
|
| if __name__ == "__main__": |
| demo.launch(css=CSS) |
