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()