File size: 2,024 Bytes
d076a57 | 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 | import os
import torch
import pandas as pd
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 = []
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:
xmin, ymin, xmax, ymax = map(float, box)
width = xmax - xmin
height = ymax - ymin
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)
if __name__ == "__main__":
TEST_IMAGE_PATH = r"rf-detr\dataset\test"
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)
|