Spaces:
Sleeping
Sleeping
| """ | |
| RT-DETR Vehicle Detector | |
| Wrapper class untuk inference menggunakan model RT-DETR yang sudah di-fine-tune | |
| pada dataset kendaraan. Nama kelas dibaca langsung dari model, jadi kompatibel | |
| dengan dataset apapun yang dipakai saat training. | |
| """ | |
| import torch | |
| from ultralytics import RTDETR | |
| from pathlib import Path | |
| class VehicleDetector: | |
| """ | |
| Wrapper untuk RT-DETR model inference. | |
| Menangani loading model, inference per frame, dan filtering | |
| berdasarkan confidence threshold. | |
| """ | |
| def __init__(self, model_path="models/best.pt", confidence=0.5): | |
| """ | |
| Args: | |
| model_path: path ke file model .pt | |
| confidence: minimum confidence score untuk deteksi | |
| """ | |
| self.model_path = Path(model_path) | |
| self.confidence = confidence | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.model = None | |
| self.class_names = {} | |
| self._load_model() | |
| def _load_model(self): | |
| """Load RT-DETR model dari file dan baca class names.""" | |
| if not self.model_path.exists(): | |
| raise FileNotFoundError( | |
| f"Model tidak ditemukan di {self.model_path}. " | |
| f"Pastikan file best.pt sudah ada di folder models/" | |
| ) | |
| self.model = RTDETR(str(self.model_path)) | |
| self.model.to(self.device) | |
| # baca class names langsung dari model | |
| # model.names biasanya berupa dict {0: 'class0', 1: 'class1', ...} | |
| if hasattr(self.model, "names") and self.model.names: | |
| self.class_names = self.model.names | |
| else: | |
| # fallback kalau names tidak ada | |
| self.class_names = {0: "Vehicle"} | |
| print(f"Model loaded pada device: {self.device}") | |
| print(f"Kelas terdeteksi: {self.class_names}") | |
| def detect(self, frame): | |
| """ | |
| Jalankan deteksi pada satu frame. | |
| Args: | |
| frame: numpy array (BGR image dari OpenCV) | |
| Returns: | |
| list of dict, masing-masing berisi: | |
| - bbox: [x1, y1, x2, y2] | |
| - confidence: float | |
| - class_id: int | |
| - class_name: str | |
| """ | |
| results = self.model.predict( | |
| frame, | |
| conf=self.confidence, | |
| device=self.device, | |
| verbose=False | |
| ) | |
| detections = [] | |
| for result in results: | |
| boxes = result.boxes | |
| if boxes is None or len(boxes) == 0: | |
| continue | |
| for i in range(len(boxes)): | |
| bbox = boxes.xyxy[i].cpu().numpy().tolist() | |
| conf = float(boxes.conf[i].cpu().numpy()) | |
| cls_id = int(boxes.cls[i].cpu().numpy()) | |
| # ambil nama kelas, fallback ke "Unknown" kalau index diluar mapping | |
| cls_name = self.class_names.get(cls_id, "Unknown") | |
| detections.append({ | |
| "bbox": bbox, | |
| "confidence": conf, | |
| "class_id": cls_id, | |
| "class_name": cls_name | |
| }) | |
| return detections | |
| def set_confidence(self, confidence): | |
| """Update confidence threshold.""" | |
| self.confidence = confidence | |
| def get_model_info(self): | |
| """Return informasi model untuk ditampilkan di UI.""" | |
| info = { | |
| "model_path": str(self.model_path), | |
| "device": self.device, | |
| "confidence_threshold": self.confidence, | |
| "num_classes": len(self.class_names), | |
| "classes": list(self.class_names.values()) | |
| } | |
| return info | |