| import cv2, json, os, sys, uuid, threading, time, base64 |
| import numpy as np |
| from ultralytics import YOLO |
| from collections import defaultdict |
|
|
| YOLO_MODEL = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'best.pt') |
| TRACKER_CONFIG = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bytetrack.yaml') |
| CONF_THRESH = 0.25 |
| IOU_THRESH = 0.45 |
| MAX_ABSENT_FRAMES = 15 |
| SIMILARITY_IOU = 0.35 |
|
|
| _model_lock = threading.Lock() |
| _model_cache = None |
|
|
|
|
| def _get_model(): |
| global _model_cache |
| if _model_cache is None: |
| _model_cache = YOLO(YOLO_MODEL) |
| return _model_cache |
|
|
|
|
| def _ensure_tracker_config(): |
| path = TRACKER_CONFIG |
| if not os.path.exists(path): |
| with open(path, 'w') as f: |
| f.write("""# ByteTrack |
| tracker_type: bytetrack |
| track_high_thresh: 0.25 |
| track_low_thresh: 0.1 |
| new_track_thresh: 0.6 |
| track_buffer: 30 |
| match_thresh: 0.8 |
| fuse_score: True |
| # BoT-SORT (uncomment to use instead) |
| # tracker_type: botsort |
| # track_high_thresh: 0.25 |
| # track_low_thresh: 0.1 |
| # new_track_thresh: 0.5 |
| # track_buffer: 50 |
| # match_thresh: 0.8 |
| # fuse_score: True |
| # gmc_method: sparseOptFlow |
| """) |
| return path |
|
|
|
|
| class AppleTracker: |
| def __init__(self, tracker_type='bytetrack'): |
| self.model = _get_model() |
| self.tracker_type = tracker_type |
| self.tracker_cfg = _ensure_tracker_config() |
| self.apple_tracks = {} |
| self.next_id = 1 |
| self.frame_count = 0 |
| self.completed_tracks = set() |
|
|
| def _get_label(self, cls_id): |
| APPLE_CLASSES = {0: 'apple'} |
| return APPLE_CLASSES.get(int(cls_id), 'apple') |
|
|
| def process_frame(self, frame): |
| h, w = frame.shape[:2] |
| self.frame_count += 1 |
|
|
| results = self.model.track( |
| frame, |
| persist=True, |
| conf=CONF_THRESH, |
| iou=IOU_THRESH, |
| tracker=self.tracker_cfg, |
| verbose=False, |
| device='cpu' |
| ) |
|
|
| current_ids = set() |
| detections = [] |
| if results and len(results) > 0: |
| r = results[0] |
| if r.boxes is not None and r.boxes.id is not None: |
| for i, box in enumerate(r.boxes): |
| track_id = int(box.id[i].item()) |
| conf = float(box.conf[i].item()) |
| cls = int(box.cls[i].item()) |
| x1, y1, x2, y2 = box.xyxy[i].tolist() |
| x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) |
| bw, bh = x2 - x1, y2 - y1 |
| cx, cy = x1 + bw // 2, y1 + bh // 2 |
| radius = max(bw, bh) // 2 |
|
|
| current_ids.add(track_id) |
| detections.append({ |
| 'frame': self.frame_count, |
| 'track_id': track_id, |
| 'bbox': [x1, y1, x2, y2], |
| 'center': [cx, cy], |
| 'radius': radius, |
| 'confidence': round(conf * 100, 1), |
| }) |
|
|
| if track_id not in self.apple_tracks: |
| self.apple_tracks[track_id] = { |
| 'track_id': track_id, |
| 'first_frame': self.frame_count, |
| 'last_frame': self.frame_count, |
| 'total_frames': 1, |
| 'max_confidence': conf, |
| 'bboxes': [detections[-1]], |
| 'absent_frames': 0, |
| } |
| else: |
| t = self.apple_tracks[track_id] |
| t['last_frame'] = self.frame_count |
| t['total_frames'] += 1 |
| t['absent_frames'] = 0 |
| if conf > t['max_confidence']: |
| t['max_confidence'] = conf |
| if len(t['bboxes']) < 60: |
| t['bboxes'].append(detections[-1]) |
| else: |
| t['bboxes'][-1] = detections[-1] |
|
|
| for tid in list(self.apple_tracks.keys()): |
| if tid not in current_ids: |
| t = self.apple_tracks[tid] |
| t['absent_frames'] += 1 |
| if t['absent_frames'] >= MAX_ABSENT_FRAMES: |
| if tid not in self.completed_tracks: |
| self.completed_tracks.add(tid) |
|
|
| return detections, current_ids |
|
|
| def get_summary(self): |
| active = {tid: t for tid, t in self.apple_tracks.items() |
| if tid not in self.completed_tracks and t['absent_frames'] < 5} |
| completed_ids = self.completed_tracks | { |
| tid for tid, t in self.apple_tracks.items() |
| if t['absent_frames'] >= 5 |
| } |
| valid_tracks = [t for tid, t in self.apple_tracks.items() if tid in completed_ids] |
|
|
| total_unique = len(valid_tracks) |
| if total_unique == 0: |
| total_unique = len(self.apple_tracks) |
|
|
| avg_conf = np.mean([t['max_confidence'] for t in self.apple_tracks.values()]) * 100 if self.apple_tracks else 0 |
|
|
| return { |
| 'total_unique_apples': total_unique, |
| 'active_tracks': len(active), |
| 'completed_tracks': len(completed_ids), |
| 'total_frames_processed': self.frame_count, |
| 'average_confidence': round(avg_conf, 1), |
| 'tracks': [ |
| { |
| 'track_id': t['track_id'], |
| 'first_frame': t['first_frame'], |
| 'last_frame': t['last_frame'], |
| 'total_frames': t['total_frames'], |
| 'max_confidence': round(t['max_confidence'] * 100, 1), |
| 'status': 'completed' if tid in completed_ids else 'active', |
| } |
| for tid, t in self.apple_tracks.items() |
| ], |
| } |
|
|
| def reset(self): |
| self.apple_tracks.clear() |
| self.completed_tracks.clear() |
| self.next_id = 1 |
| self.frame_count = 0 |
|
|
|
|
| def create_tracked_annotation(frame, detections, active_ids, completed_count): |
| annotated = frame.copy() |
| h, w = frame.shape[:2] |
| font = cv2.FONT_HERSHEY_SIMPLEX |
| font_scale = max(0.5, min(h, w) / 1500.0) |
| thickness = max(2, int(font_scale * 2.5)) |
| line_w = max(3, int(min(h, w) / 300)) |
|
|
| for det in detections: |
| tid = det['track_id'] |
| x1, y1, x2, y2 = det['bbox'] |
| color = (0, 200, 0) if tid in active_ids else (100, 100, 100) |
| cv2.rectangle(annotated, (x1, y1), (x2, y2), color, line_w) |
| label = f"#{tid}" |
| (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness) |
| pad = 4 |
| ly = y1 - 8 if y1 - th - pad > 0 else y2 + th + pad |
| cv2.rectangle(annotated, (x1, ly - th - pad), (x1 + tw + pad * 2, ly + pad), (0, 0, 0), -1) |
| cv2.putText(annotated, label, (x1 + pad, ly), font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA) |
|
|
| status = f"Frame: {detections[0]['frame'] if detections else 0} | Active: {len(active_ids)} | Total unique: {completed_count}" |
| (stw, sth), _ = cv2.getTextSize(status, font, font_scale * 1.1, thickness) |
| bar_h = sth + 16 |
| cv2.rectangle(annotated, (0, 0), (stw + 24, bar_h), (0, 0, 0), -1) |
| cv2.putText(annotated, status, (12, sth + 8), font, font_scale * 1.1, (0, 255, 0), thickness, cv2.LINE_AA) |
|
|
| return annotated |
|
|
|
|
| def process_video_file(video_path, max_frames=0, progress_callback=None): |
| cap = cv2.VideoCapture(video_path) |
| if not cap.isOpened(): |
| raise ValueError(f"Cannot open video: {video_path}") |
|
|
| fps = cap.get(cv2.CAP_PROP_FPS) |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAMES)) |
| if max_frames > 0: |
| total_frames = min(total_frames, max_frames) |
|
|
| tracker = AppleTracker() |
| annotated_frames = [] |
| frame_idx = 0 |
|
|
| while True: |
| if max_frames > 0 and frame_idx >= max_frames: |
| break |
| ret, frame = cap.read() |
| if not ret: |
| break |
|
|
| detections, active_ids = tracker.process_frame(frame) |
| if frame_idx % max(1, int(fps / 3)) == 0 or len(detections) > 0: |
| annotated = create_tracked_annotation( |
| frame, detections, active_ids, |
| len(tracker.completed_tracks) |
| ) |
| _, buffer = cv2.imencode('.jpg', annotated, [cv2.IMWRITE_JPEG_QUALITY, 85]) |
| annotated_b64 = base64.b64encode(buffer).decode('utf-8') |
| annotated_frames.append({ |
| 'frame': frame_idx + 1, |
| 'active_tracks': len(active_ids), |
| 'detections': len(detections), |
| 'annotated_base64': annotated_b64, |
| }) |
|
|
| if len(annotated_frames) > 20: |
| annotated_frames.pop(0) |
|
|
| frame_idx += 1 |
| if progress_callback and frame_idx % max(1, int(total_frames / 10)) == 0: |
| progress_callback(int(frame_idx / total_frames * 100)) |
|
|
| cap.release() |
| summary = tracker.get_summary() |
|
|
| return { |
| 'summary': summary, |
| 'fps': fps, |
| 'total_frames': frame_idx, |
| 'annotated_previews': annotated_frames, |
| } |
|
|
