fu4ll/Vehicle_Detection

1from ultralytics import YOLO2import cv23import base644import subprocess5import os6 7model = YOLO("yolov8n.pt")8car_model = YOLO("best.pt")9 10VEHICLE_CLASSES = {"car", "truck", "bus", "motorcycle", "bicycle", "airplane", "boat"}11 12COLORS = [13    (255, 99, 99),  (99, 255, 99),  (99, 99, 255),14    (255, 199, 99), (99, 255, 255), (255, 99, 255),15    (180, 255, 99), (99, 180, 255), (255, 99, 180),16    (200, 200, 99),17]18 19def get_color(track_id: int):20    return COLORS[track_id % len(COLORS)]21 22def draw_box(frame, x1, y1, x2, y2, label, color):23    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)24    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)25    cv2.rectangle(frame, (x1, y1 - th - 8), (x1 + tw + 6, y1), color, -1)26    cv2.putText(frame, label, (x1 + 3, y1 - 4),27        cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 1, cv2.LINE_AA)28 29def classify_car(img, x1, y1, x2, y2):30    crop = img[y1:y2, x1:x2]31    if crop.size == 0:32        return "unknown"33    results = car_model(crop, verbose=False)34    for r in results:35        for box in r.boxes:36            return car_model.names[int(box.cls)]37    return "unknown"38 39def detect(image_path: str):40    results = model(image_path)41    detections = []42    img = cv2.imread(image_path)43 44    for idx, r in enumerate(results):45        for i, box in enumerate(r.boxes):46            cls = model.names[int(box.cls)]47            if cls not in VEHICLE_CLASSES:48                continue49            conf = round(float(box.conf), 2)50            x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]51            track_id = i + 152            color = get_color(track_id)53 54            car_type = classify_car(img, x1, y1, x2, y2)55 56            detections.append({57                "id": track_id,58                "class": cls,59                "car_model": car_type,60                "confidence": conf,61                "bbox": [x1, y1, x2, y2],62                "color": f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}"63            })64 65            label = f"#{track_id} {car_type} {conf}"66            draw_box(img, x1, y1, x2, y2, label, color)67 68    _, buffer = cv2.imencode(".jpg", img)69    img_base64 = base64.b64encode(buffer).decode("utf-8")70    return detections, img_base6471 72 73def detect_video(video_path: str, output_path: str):74    cap = cv2.VideoCapture(video_path)75    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))76    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))77    fps    = cap.get(cv2.CAP_PROP_FPS)78 79    fourcc = cv2.VideoWriter_fourcc(*"avc1")80    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))81 82    unique_ids = set()83    class_counts = {}84    id_colors = {}85 86    while cap.isOpened():87        ret, frame = cap.read()88        if not ret:89            break90 91        results = model.track(frame, verbose=False, persist=True)92 93        for r in results:94            if r.boxes.id is None:95                continue96            for box, track_id in zip(r.boxes, r.boxes.id):97                cls = model.names[int(box.cls)]98                if cls not in VEHICLE_CLASSES:99                    continue100 101                tid = int(track_id)102                unique_ids.add(tid)103                color = get_color(tid)104                id_colors[tid] = f"#{color[0]:02x}{color[1]:02x}{color[2]:02x}"105 106                conf = round(float(box.conf), 2)107                x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]108 109                car_type = classify_car(frame, x1, y1, x2, y2)110 111                if tid not in class_counts:112                    class_counts[tid] = car_type113 114                draw_box(frame, x1, y1, x2, y2,115                    f"#{tid} {car_type} {conf}", color)116 117        out.write(frame)118 119    cap.release()120    out.release()121 122    converted_path = output_path.replace(".mp4", "_web.mp4")123    subprocess.run([124        "ffmpeg", "-y", "-i", output_path,125        "-vcodec", "libx264", "-acodec", "aac",126        converted_path127    ], capture_output=True)128 129    stats = {}130    for tid, cls in class_counts.items():131        stats[cls] = stats.get(cls, 0) + 1132 133    objects = [134        {"id": tid, "class": cls, "color": id_colors.get(tid, "#ffffff")}135        for tid, cls in class_counts.items()136    ]137 138    return {139        "unique_vehicles": len(unique_ids),140        "by_class": stats,141        "objects": objects142    }, converted_path