| import os |
| import torch |
| import pandas as pd |
| from PIL import Image |
| from rfdetr import RFDETRBase |
|
|
|
|
| def run_inference(model, image_path, conf_threshold, save_path): |
|
|
| test_images = sorted([ |
| f for f in os.listdir(image_path) |
| if f.lower().endswith((".jpg", ".jpeg", ".png")) |
| ]) |
|
|
| bboxes = [] |
| category_ids = [] |
| test_images_names = [] |
|
|
| for image_name in test_images: |
| test_images_names.append(image_name) |
| image_file = os.path.join(image_path, image_name) |
|
|
| bbox = [] |
| category_id = [] |
|
|
| |
| with Image.open(image_file) as img: |
| img_w, img_h = img.size |
|
|
| preds = model.predict(image_file) |
|
|
| if preds is not None and preds.xyxy is not None and len(preds.xyxy) > 0: |
| for box, score, label in zip( |
| preds.xyxy, |
| preds.confidence, |
| preds.class_id |
| ): |
| score = float(score) |
| if score < conf_threshold: |
| continue |
|
|
| xmin, ymin, xmax, ymax = map(float, box) |
|
|
| |
| xmin = max(0.0, xmin) |
| ymin = max(0.0, ymin) |
| xmax = min(float(img_w), xmax) |
| ymax = min(float(img_h), ymax) |
|
|
| width = xmax - xmin |
| height = ymax - ymin |
|
|
| if width <= 0 or height <= 0: |
| continue |
|
|
| bbox.append([xmin, ymin, width, height]) |
| category_id.append(int(label)) |
|
|
| bboxes.append(bbox) |
| category_ids.append(category_id) |
|
|
| df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"]) |
|
|
| for i in range(len(test_images_names)): |
| new_row = pd.DataFrame({ |
| "file_name": test_images_names[i], |
| "bbox": str(bboxes[i]), |
| "category_id": str(category_ids[i]) |
| }, index=[0]) |
|
|
| df_predictions = pd.concat([df_predictions, new_row], ignore_index=True) |
|
|
| df_predictions.to_csv(save_path, index=False) |
| print(f"Submission saved to {save_path}") |
|
|
|
|
| if __name__ == "__main__": |
|
|
| TEST_IMAGE_PATH = "/tmp/data/test_images" |
| SUBMISSION_SAVE_PATH = "submission.csv" |
| CONF_THRESHOLD = 0.30 |
|
|
| model = RFDETRBase( |
| checkpoint_path="checkpoint_best_total.pth", |
| device="cuda" if torch.cuda.is_available() else "cpu" |
| ) |
|
|
| run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH) |
|
|
