yolov8-streamlit-test / object_detection_script.py
ethanrom's picture
initial
5b5b64e
import os
import cv2
from ultralytics import YOLO
from tqdm import tqdm
model = YOLO('yolov8n.pt')
def detect_cars_humans(image_path):
image_cv = cv2.imread(image_path)
# Perform object detection
results = model(
image_cv,
classes=[0, 2, 7, 3, 5],
iou=0.7,
conf=0.65,
show_labels=False,
show_conf=False,
boxes=True
)
if len(results[0].boxes.xyxy) == 0:
return
# Create the destination folder if it doesn't exist
os.makedirs(r"D:\ascii\car_person\testing\Test_Purpose\Test1_results", exist_ok=True)
# Save the annotated image in the results folder
annotated_image_path = os.path.join(r"D:\ascii\car_person\testing\Test_Purpose\Test1_results", os.path.basename(image_path))
cv2.imwrite(annotated_image_path, results[0].plot())
source_folder = r"D:\ascii\car_person\testing\Test_Purpose\Test1"
image_files = [f for f in os.listdir(source_folder) if f.endswith(".png") or f.endswith(".jpg")]
with tqdm(total=len(image_files), desc='Processing Images') as pbar:
for filename in image_files:
image_path = os.path.join(source_folder, filename)
detect_cars_humans(image_path)
pbar.update(1)