File size: 3,015 Bytes
333330e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import os
import pandas as pd
from ultralytics import YOLO
def run_inference(model, image_path, conf_threshold, save_path):
test_images = os.listdir(image_path)
test_images.sort()
bboxes = []
category_ids = []
test_images_names = []
# Iterate through images for inference
for image_name in test_images:
# Skip any non-image files if they exist in the directory
if not image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
continue
full_image_path = os.path.join(image_path, image_name)
current_image_bboxes = []
current_image_category_ids = []
# Perform inference
results = model(full_image_path)
# Process results
# results is a list of Results objects, one for each image
# Since we pass one image at a time, results[0] is the relevant object
for pred in results[0].boxes:
# Bounding box in xyxy format, confidence, class_id
xmin, ymin, xmax, ymax = pred.xyxy[0].tolist()
conf = pred.conf.item()
class_id = int(pred.cls.item())
if conf >= conf_threshold:
width = xmax - xmin
height = ymax - ymin
current_image_bboxes.append([xmin, ymin, width, height])
current_image_category_ids.append(class_id)
test_images_names.append(image_name)
bboxes.append(current_image_bboxes)
category_ids.append(current_image_category_ids)
# Create DataFrame for predictions
df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
for i in range(len(test_images_names)):
file_name = test_images_names[i]
new_row = pd.DataFrame({"file_name": file_name,
"bbox": str(bboxes[i]),
"category_id": str(category_ids[i]),
}, index=[0])
df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
# Ensure the save directory exists
os.makedirs(os.path.dirname(save_path), exist_ok=True)
df_predictions.to_csv(save_path, index=False)
print(f"Inference results saved to: {save_path}")
if __name__ == "__main__":
# Define paths
# You might need to change TEST_IMAGE_PATH to where your actual test images are stored
current_directory = os.path.dirname(os.path.abspath(__file__))
TEST_IMAGE_PATH = "/tmp/data/test_images"
SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv")
# Path to your trained model weights
MODEL_WEIGHTS_PATH = os.path.join(current_directory, "best.pt")
CONF_THRESHOLD = 0.30 # Confidence threshold for predictions
# Load the YOLO model
model = YOLO(MODEL_WEIGHTS_PATH) # Using ultralytics.YOLO for loading
# Run inference
run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)
|