app.py

import os
import uuid
import pandas as pd
import cv2
from ultralytics import YOLO

# Initialize the YOLO model
model = YOLO('best.pt')

# Define directories
image_dir = 'datasetsw/valid/ripe'
output_dir = 'datasetsw/valid/test_ripe'

# List all image files in the directory
image_files = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]

# Initialize an empty list to store the results
results_list = []

# Process each image
for image_file in image_files:
    results = model(image_file)
    annotated_frame = results[0].plot()

    total_objects = len(results[0].boxes)
    labels = results[0].boxes.cls.tolist()
    matang_count = labels.count(0)
    mentah_count = labels.count(1)

    # Generate a unique filename and save the annotated image
    unique_id = str(uuid.uuid4())
    combined_filename = f"{unique_id}.jpg"
    output_path = os.path.join(output_dir, combined_filename)
    cv2.imwrite(output_path, annotated_frame)

    # Append the results to the list
    results_list.append({
        'image_file': image_file,
        'total_objects': total_objects,
        'matang_count': matang_count,
        'mentah_count': mentah_count
    })

# Convert the results list to a dataframe
results_df = pd.DataFrame(results_list)

# Save the dataframe to a CSV file
results_df.to_csv(os.path.join(output_dir, 'detection_results.csv'), index=False)

cv2.destroyAllWindows()