CalorieTrackerAI-backend / testing_yolo.py
cosmin1769's picture
Upload 124 files
21ff07a verified
from ultralytics import YOLO
from pathlib import Path
import cv2
import random
model = YOLO('runs/detect/allinone_yolov8n_v1/weights/best.pt')
val_images = Path('data/Dataset2_rboflow/valid/images')
val_labels = Path('data/Dataset2_rboflow/valid/labels')
output_dir = Path('test_comparison')
output_dir.mkdir(exist_ok=True)
all_images = list(val_images.glob('*.jpg'))
test_images = random.sample(all_images, min(10, len(all_images)))
print(f"Testing pe {len(test_images)} imagini\n")
for img_path in test_images:
results = model.predict(img_path, conf=0.25, imgsz=640, verbose=False)
n_detections = len(results[0].boxes)
label_path = val_labels / f"{img_path.stem}.txt"
n_true = 0
if label_path.exists():
with open(label_path, 'r') as f:
n_true = len(f.readlines())
result_img = results[0].plot()
output_path = output_dir / img_path.name
cv2.imwrite(str(output_path), result_img)
print(f" {img_path.name}")
print(f" TRUE objects: {n_true}")
print(f" PREDICTED: {n_detections}")
print(f" Salvat in: {output_path}\n")