mustafa2ak's picture
Update detector.py
c6b8113 verified
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np
class RoadsideDetector:
def __init__(self, model_path="yol_kenari_tespit_best.pt"):
self.model = YOLO(model_path)
# Türkçe sınıf isimleri
self.class_names = {
0: "çöp",
1: "çöp kutusu",
2: "taşan çöp",
}
# Seviye hesaplamasına giren sınıflar
self.severity_classes = ["çöp", "taşan çöp"]
def detect_items(self, image):
results = self.model(image, verbose=False)
detections = []
for r in results:
boxes = r.boxes.xyxy.cpu().numpy()
confs = r.boxes.conf.cpu().numpy()
classes = r.boxes.cls.cpu().numpy()
for box, conf, cls_id in zip(boxes, confs, classes):
cls_name = self.class_names.get(int(cls_id), "Bilinmeyen")
detections.append({
"class_name": cls_name,
"conf": float(conf),
"box": box.tolist(),
})
return detections
def draw_boxes(self, image, detections):
"""Türkçe karakter desteği ile bounding box çizimi"""
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
draw = ImageDraw.Draw(image)
try:
font = ImageFont.truetype("DejaVuSans.ttf", 28)
except:
font = ImageFont.load_default()
for det in detections:
box = det["box"]
cls_name = det["class_name"]
conf = det["conf"]
x1, y1, x2, y2 = box
draw.rectangle([x1, y1, x2, y2], outline="lime", width=4)
label = f"{cls_name} ({conf:.2f})"
# ✅ Yeni Pillow sürümü için textsize yerine textbbox kullanalım
bbox = draw.textbbox((0, 0), label, font=font)
text_w = bbox[2] - bbox[0]
text_h = bbox[3] - bbox[1]
draw.rectangle([x1, y1 - text_h - 4, x1 + text_w + 4, y1], fill="black")
draw.text((x1 + 2, y1 - text_h - 2), label, font=font, fill="yellow")
return image
def calculate_severity(self, detections):
"""Sadece çöp ve taşan çöp seviye hesabına girer"""
severity_count = sum(1 for d in detections if d["class_name"] in self.severity_classes)
return severity_count