| from pathlib import Path |
| import csv |
| import uuid |
|
|
| import cv2 |
| import gradio as gr |
| from sahi import AutoDetectionModel |
| from sahi.predict import get_sliced_prediction |
| from ultralytics import YOLO |
|
|
|
|
| APP_DIR = Path(__file__).resolve().parent |
| MODEL_PATH = APP_DIR / "best.onnx" |
| OUTPUTS_DIR = APP_DIR / "outputs" |
| TILE_SIZE = 640 |
| OVERLAP_RATIO = 0.25 |
| CONFIDENCE_THRESHOLD = 0.5 |
| NMS_MATCH_THRESHOLD = 0.3 |
| MAX_TRACK_MISSES = 4 |
| MIN_TRACK_FRAMES = 3 |
| MOTION_THRESHOLD_PX = 2.0 |
| MAX_CENTER_DISTANCE_PX = 120.0 |
| MIN_IOU_FOR_MATCH = 0.05 |
| CLASS_COLORS = { |
| "predator": (255, 99, 71), |
| "prey": (0, 200, 83), |
| } |
|
|
| base_model = YOLO(str(MODEL_PATH), task="detect") |
| detection_model = AutoDetectionModel.from_pretrained( |
| model_type="ultralytics", |
| model=base_model, |
| model_path=str(MODEL_PATH), |
| confidence_threshold=CONFIDENCE_THRESHOLD, |
| device="cpu", |
| load_at_init=True, |
| ) |
|
|
|
|
| def intersection_over_smaller(box_a, box_b): |
| ax1, ay1, ax2, ay2 = box_a |
| bx1, by1, bx2, by2 = box_b |
|
|
| inter_x1 = max(ax1, bx1) |
| inter_y1 = max(ay1, by1) |
| inter_x2 = min(ax2, bx2) |
| inter_y2 = min(ay2, by2) |
|
|
| inter_w = max(0.0, inter_x2 - inter_x1) |
| inter_h = max(0.0, inter_y2 - inter_y1) |
| intersection = inter_w * inter_h |
| if intersection <= 0: |
| return 0.0 |
|
|
| area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1) |
| area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1) |
| smaller_area = min(area_a, area_b) |
| if smaller_area <= 0: |
| return 0.0 |
|
|
| return intersection / smaller_area |
|
|
|
|
| def bbox_iou(box_a, box_b): |
| ax1, ay1, ax2, ay2 = box_a |
| bx1, by1, bx2, by2 = box_b |
|
|
| inter_x1 = max(ax1, bx1) |
| inter_y1 = max(ay1, by1) |
| inter_x2 = min(ax2, bx2) |
| inter_y2 = min(ay2, by2) |
|
|
| inter_w = max(0.0, inter_x2 - inter_x1) |
| inter_h = max(0.0, inter_y2 - inter_y1) |
| intersection = inter_w * inter_h |
| if intersection <= 0: |
| return 0.0 |
|
|
| area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1) |
| area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1) |
| union = area_a + area_b - intersection |
| if union <= 0: |
| return 0.0 |
|
|
| return intersection / union |
|
|
|
|
| def normalized_center_distance(box_a, box_b): |
| ax1, ay1, ax2, ay2 = box_a |
| bx1, by1, bx2, by2 = box_b |
|
|
| center_a = ((ax1 + ax2) / 2.0, (ay1 + ay2) / 2.0) |
| center_b = ((bx1 + bx2) / 2.0, (by1 + by2) / 2.0) |
| center_distance = ((center_a[0] - center_b[0]) ** 2 + (center_a[1] - center_b[1]) ** 2) ** 0.5 |
|
|
| width_a = max(0.0, ax2 - ax1) |
| height_a = max(0.0, ay2 - ay1) |
| width_b = max(0.0, bx2 - bx1) |
| height_b = max(0.0, by2 - by1) |
| smaller_diagonal = min((width_a ** 2 + height_a ** 2) ** 0.5, (width_b ** 2 + height_b ** 2) ** 0.5) |
|
|
| if smaller_diagonal <= 0: |
| return float("inf") |
|
|
| return center_distance / smaller_diagonal |
|
|
|
|
| def resolve_cross_class_duplicates(predictions): |
| if len(predictions) < 2: |
| return predictions |
|
|
| duplicate_ios_threshold = 0.85 |
| center_distance_ratio_threshold = 0.35 |
| kept_predictions = [] |
|
|
| for pred in sorted( |
| predictions, |
| key=lambda item: float(getattr(item.score, "value", 0.0)), |
| reverse=True, |
| ): |
| should_skip = False |
| current_bbox = pred.bbox.to_xyxy() |
| current_class = pred.category.name |
|
|
| for kept in kept_predictions: |
| if kept.category.name == current_class: |
| continue |
|
|
| kept_bbox = kept.bbox.to_xyxy() |
| ios = intersection_over_smaller(current_bbox, kept_bbox) |
| center_distance_ratio = normalized_center_distance(current_bbox, kept_bbox) |
|
|
| if ios >= duplicate_ios_threshold and center_distance_ratio <= center_distance_ratio_threshold: |
| should_skip = True |
| break |
|
|
| if not should_skip: |
| kept_predictions.append(pred) |
|
|
| return kept_predictions |
|
|
|
|
| def run_sliced_detection(image): |
| result = get_sliced_prediction( |
| image, |
| detection_model, |
| slice_height=TILE_SIZE, |
| slice_width=TILE_SIZE, |
| overlap_height_ratio=OVERLAP_RATIO, |
| overlap_width_ratio=OVERLAP_RATIO, |
| postprocess_type="GREEDYNMM", |
| postprocess_match_metric="IOS", |
| postprocess_match_threshold=NMS_MATCH_THRESHOLD, |
| postprocess_class_agnostic=False, |
| perform_standard_pred=False, |
| verbose=0, |
| ) |
| return resolve_cross_class_duplicates(result.object_prediction_list) |
|
|
|
|
| def prediction_to_detection(pred): |
| x1, y1, x2, y2 = [float(v) for v in pred.bbox.to_xyxy()] |
| center_x = (x1 + x2) / 2.0 |
| center_y = (y1 + y2) / 2.0 |
| return { |
| "bbox": [x1, y1, x2, y2], |
| "class_name": pred.category.name, |
| "score": float(getattr(pred.score, "value", 0.0)), |
| "center": [center_x, center_y], |
| } |
|
|
|
|
| def render_image(image_rgb, detections): |
| rendered = image_rgb.copy() |
|
|
| for det in detections: |
| x1, y1, x2, y2 = [int(v) for v in det["bbox"]] |
| label = det["class_name"] |
| score = det["score"] |
| track_id = det.get("track_id") |
| color = CLASS_COLORS.get(label, (255, 193, 7)) |
|
|
| cv2.rectangle(rendered, (x1, y1), (x2, y2), color, 3) |
|
|
| text = f"{label} {score:.2f}" |
| if track_id is not None: |
| text = f"ID {track_id} | {text}" |
|
|
| (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) |
| text_y = max(y1 - 10, text_height + 8) |
| cv2.rectangle( |
| rendered, |
| (x1, text_y - text_height - 8), |
| (x1 + text_width + 10, text_y), |
| color, |
| -1, |
| ) |
| cv2.putText( |
| rendered, |
| text, |
| (x1 + 5, text_y - 5), |
| cv2.FONT_HERSHEY_SIMPLEX, |
| 0.7, |
| (255, 255, 255), |
| 2, |
| cv2.LINE_AA, |
| ) |
|
|
| return rendered |
|
|
|
|
| def detect_images(files): |
| if not files: |
| return [] |
|
|
| outputs = [] |
| for file_path in files: |
| image_bgr = cv2.imread(str(file_path)) |
| if image_bgr is None: |
| continue |
| image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) |
| detections = [prediction_to_detection(pred) for pred in run_sliced_detection(image_rgb)] |
| outputs.append(render_image(image_rgb, detections)) |
| return outputs |
|
|
|
|
| def build_track(track_id, detection, frame_idx): |
| return { |
| "track_id": track_id, |
| "class_name": detection["class_name"], |
| "bbox": detection["bbox"], |
| "center": detection["center"], |
| "last_frame_idx": frame_idx, |
| "frames_seen": 1, |
| "misses": 0, |
| "total_distance_px": 0.0, |
| "moving_frames": 0, |
| "sum_confidence": detection["score"], |
| } |
|
|
|
|
| def assign_tracks(detections, tracks, next_track_id, frame_idx): |
| matched_detections = set() |
| matched_tracks = set() |
| candidate_pairs = [] |
|
|
| for det_idx, detection in enumerate(detections): |
| for track_id, track in tracks.items(): |
| if track["class_name"] != detection["class_name"]: |
| continue |
| if track["misses"] > MAX_TRACK_MISSES: |
| continue |
|
|
| prev_center = track["center"] |
| curr_center = detection["center"] |
| center_distance = ((prev_center[0] - curr_center[0]) ** 2 + (prev_center[1] - curr_center[1]) ** 2) ** 0.5 |
| iou = bbox_iou(track["bbox"], detection["bbox"]) |
| dynamic_threshold = max( |
| MAX_CENTER_DISTANCE_PX, |
| 0.75 * min( |
| ((track["bbox"][2] - track["bbox"][0]) ** 2 + (track["bbox"][3] - track["bbox"][1]) ** 2) ** 0.5, |
| ((detection["bbox"][2] - detection["bbox"][0]) ** 2 + (detection["bbox"][3] - detection["bbox"][1]) ** 2) ** 0.5, |
| ), |
| ) |
|
|
| if center_distance <= dynamic_threshold or iou >= MIN_IOU_FOR_MATCH: |
| candidate_pairs.append((-iou, center_distance, track_id, det_idx)) |
|
|
| for _, center_distance, track_id, det_idx in sorted(candidate_pairs): |
| if track_id in matched_tracks or det_idx in matched_detections: |
| continue |
|
|
| detection = detections[det_idx] |
| track = tracks[track_id] |
| track["total_distance_px"] += center_distance |
| track["moving_frames"] += int(center_distance >= MOTION_THRESHOLD_PX) |
| track["bbox"] = detection["bbox"] |
| track["center"] = detection["center"] |
| track["last_frame_idx"] = frame_idx |
| track["frames_seen"] += 1 |
| track["misses"] = 0 |
| track["sum_confidence"] += detection["score"] |
| detection["track_id"] = track_id |
| matched_tracks.add(track_id) |
| matched_detections.add(det_idx) |
|
|
| for track_id, track in tracks.items(): |
| if track_id not in matched_tracks: |
| track["misses"] += 1 |
|
|
| for det_idx, detection in enumerate(detections): |
| if det_idx in matched_detections: |
| continue |
|
|
| track_id = next_track_id |
| next_track_id += 1 |
| detection["track_id"] = track_id |
| tracks[track_id] = build_track(track_id, detection, frame_idx) |
|
|
| return next_track_id |
|
|
|
|
| def write_tracks_csv(path, rows): |
| with path.open("w", newline="", encoding="utf-8") as handle: |
| writer = csv.DictWriter( |
| handle, |
| fieldnames=[ |
| "frame_idx", |
| "track_id", |
| "class_name", |
| "confidence", |
| "center_x", |
| "center_y", |
| "x1", |
| "y1", |
| "x2", |
| "y2", |
| ], |
| ) |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
|
|
| def summarize_tracks(tracks, frame_count, fps): |
| confirmed_tracks = [track for track in tracks.values() if track["frames_seen"] >= MIN_TRACK_FRAMES] |
| observed_seconds = frame_count / fps if fps else 0.0 |
|
|
| summary_rows = [] |
| for track in confirmed_tracks: |
| duration_seconds = track["frames_seen"] / fps if fps else 0.0 |
| avg_speed = track["total_distance_px"] / duration_seconds if duration_seconds else 0.0 |
| movement_ratio = track["moving_frames"] / max(track["frames_seen"] - 1, 1) |
| visibility_ratio = track["frames_seen"] / max(frame_count, 1) |
| summary_rows.append( |
| { |
| "track_id": track["track_id"], |
| "class_name": track["class_name"], |
| "frames_seen": track["frames_seen"], |
| "total_distance_px": round(track["total_distance_px"], 3), |
| "avg_speed_px_per_sec": round(avg_speed, 3), |
| "movement_ratio": round(movement_ratio, 4), |
| "visibility_ratio": round(visibility_ratio, 4), |
| "avg_confidence": round(track["sum_confidence"] / track["frames_seen"], 4), |
| } |
| ) |
|
|
| moving_tracks = [ |
| row for row in summary_rows |
| if row["total_distance_px"] >= MOTION_THRESHOLD_PX * 2 |
| ] |
| avg_speed_px_per_sec = ( |
| sum(row["avg_speed_px_per_sec"] for row in summary_rows) / len(summary_rows) |
| if summary_rows else 0.0 |
| ) |
| moving_object_ratio = len(moving_tracks) / len(summary_rows) if summary_rows else 0.0 |
| observation_stability = ( |
| sum(row["visibility_ratio"] for row in summary_rows) / len(summary_rows) |
| if summary_rows else 0.0 |
| ) |
| vitality_score = avg_speed_px_per_sec * moving_object_ratio * observation_stability |
|
|
| aggregate = { |
| "frame_count": frame_count, |
| "fps": round(fps, 3), |
| "observed_seconds": round(observed_seconds, 3), |
| "confirmed_tracks": len(summary_rows), |
| "moving_tracks": len(moving_tracks), |
| "avg_speed_px_per_sec": round(avg_speed_px_per_sec, 3), |
| "moving_object_ratio": round(moving_object_ratio, 4), |
| "observation_stability": round(observation_stability, 4), |
| "vitality_score": round(vitality_score, 3), |
| } |
|
|
| return summary_rows, aggregate |
|
|
|
|
| def write_summary_csv(path, summary_rows, aggregate): |
| with path.open("w", newline="", encoding="utf-8") as handle: |
| aggregate_writer = csv.writer(handle) |
| aggregate_writer.writerow(["metric", "value"]) |
| for key, value in aggregate.items(): |
| aggregate_writer.writerow([key, value]) |
|
|
| aggregate_writer.writerow([]) |
| aggregate_writer.writerow( |
| [ |
| "track_id", |
| "class_name", |
| "frames_seen", |
| "total_distance_px", |
| "avg_speed_px_per_sec", |
| "movement_ratio", |
| "visibility_ratio", |
| "avg_confidence", |
| ] |
| ) |
| for row in summary_rows: |
| aggregate_writer.writerow( |
| [ |
| row["track_id"], |
| row["class_name"], |
| row["frames_seen"], |
| row["total_distance_px"], |
| row["avg_speed_px_per_sec"], |
| row["movement_ratio"], |
| row["visibility_ratio"], |
| row["avg_confidence"], |
| ] |
| ) |
|
|
|
|
| def build_summary_markdown(aggregate, summary_rows): |
| lines = [ |
| "### Vitality Summary", |
| f"- Frames processed: {aggregate['frame_count']}", |
| f"- Observed seconds: {aggregate['observed_seconds']}", |
| f"- Confirmed tracks: {aggregate['confirmed_tracks']}", |
| f"- Moving tracks: {aggregate['moving_tracks']}", |
| f"- Average speed: {aggregate['avg_speed_px_per_sec']} px/s", |
| f"- Moving object ratio: {aggregate['moving_object_ratio']}", |
| f"- Observation stability: {aggregate['observation_stability']}", |
| f"- Vitality score: {aggregate['vitality_score']}", |
| ] |
|
|
| if summary_rows: |
| top_tracks = sorted(summary_rows, key=lambda row: row["total_distance_px"], reverse=True)[:5] |
| lines.append("") |
| lines.append("Top tracks by distance:") |
| for row in top_tracks: |
| lines.append( |
| f"- ID {row['track_id']} ({row['class_name']}): " |
| f"{row['total_distance_px']} px over {row['frames_seen']} frames" |
| ) |
|
|
| return "\n".join(lines) |
|
|
|
|
| def track_video(video_path): |
| if not video_path: |
| return None, [], "No video provided." |
|
|
| run_dir = OUTPUTS_DIR / f"video_{uuid.uuid4().hex[:8]}" |
| run_dir.mkdir(parents=True, exist_ok=True) |
| output_video = run_dir / "tracked.mp4" |
| tracks_csv = run_dir / "tracks.csv" |
| summary_csv = run_dir / "summary.csv" |
|
|
| capture = cv2.VideoCapture(str(video_path)) |
| if not capture.isOpened(): |
| return None, [], "Failed to open the uploaded video." |
|
|
| fps = capture.get(cv2.CAP_PROP_FPS) or 10.0 |
| width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
| writer = cv2.VideoWriter( |
| str(output_video), |
| cv2.VideoWriter_fourcc(*"mp4v"), |
| fps, |
| (width, height), |
| ) |
|
|
| tracks = {} |
| track_rows = [] |
| next_track_id = 1 |
| frame_idx = 0 |
|
|
| while True: |
| ok, frame_bgr = capture.read() |
| if not ok: |
| break |
|
|
| frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) |
| predictions = run_sliced_detection(frame_rgb) |
| detections = [prediction_to_detection(pred) for pred in predictions] |
| next_track_id = assign_tracks(detections, tracks, next_track_id, frame_idx) |
|
|
| for detection in detections: |
| x1, y1, x2, y2 = detection["bbox"] |
| center_x, center_y = detection["center"] |
| track_rows.append( |
| { |
| "frame_idx": frame_idx, |
| "track_id": detection["track_id"], |
| "class_name": detection["class_name"], |
| "confidence": round(detection["score"], 4), |
| "center_x": round(center_x, 3), |
| "center_y": round(center_y, 3), |
| "x1": round(x1, 3), |
| "y1": round(y1, 3), |
| "x2": round(x2, 3), |
| "y2": round(y2, 3), |
| } |
| ) |
|
|
| rendered_rgb = render_image(frame_rgb, detections) |
| rendered_bgr = cv2.cvtColor(rendered_rgb, cv2.COLOR_RGB2BGR) |
| writer.write(rendered_bgr) |
| frame_idx += 1 |
|
|
| capture.release() |
| writer.release() |
|
|
| write_tracks_csv(tracks_csv, track_rows) |
| summary_rows, aggregate = summarize_tracks(tracks, frame_idx, fps) |
| write_summary_csv(summary_csv, summary_rows, aggregate) |
| summary_markdown = build_summary_markdown(aggregate, summary_rows) |
|
|
| return str(output_video), [str(tracks_csv), str(summary_csv)], summary_markdown |
|
|
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("# Mite Detection, Tracking, and Vitality") |
|
|
| with gr.Tab("Image Detection"): |
| image_input = gr.Files(label="Upload images", file_count="multiple", type="filepath") |
| image_output = gr.Gallery(label="Tiled detections", columns=2, preview=True) |
| image_run = gr.Button("Run Detection", variant="primary") |
| image_run.click(fn=detect_images, inputs=image_input, outputs=image_output) |
|
|
| with gr.Tab("Video Tracking"): |
| video_input = gr.Video(label="Upload a video") |
| video_output = gr.Video(label="Tracked output") |
| csv_output = gr.Files(label="Tracking CSV outputs") |
| summary_output = gr.Markdown(label="Vitality summary") |
| video_run = gr.Button("Run Tracking", variant="primary") |
| video_run.click( |
| fn=track_video, |
| inputs=video_input, |
| outputs=[video_output, csv_output, summary_output], |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
