text_plain_output_1.png
from mmdet.apis import init_detector, inference_detector

from tqdm import tqdm

import mmcv

import os

import pandas as pd

import torch

def format_prediction_string(boxes, scores):
    pred_strings = []
    for j in zip(scores, boxes):
        pred_strings.append('{0:.4f} {1} {2} {3} {4}'.format(j[0], j[1][0], j[1][1], j[1][2], j[1][3]))
    return ' '.join(pred_strings)

CONFIG_FILE = './config.py'
CHECKPOINT_PATH = './model.pth'
TEST_IMG_DIR = '../input/global-wheat-detection/test'

import torch
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

config = mmcv.Config.fromfile(CONFIG_FILE)
config.model.pretrained = None
config.model.neck.rfp_backbone.pretrained = None
if False:
    img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
    config.data.test.pipeline = [dict(type='LoadImageFromFile'), dict(type='MultiScaleFlipAug', img_scale=(1024, 1024), flip=False, transforms=[dict(type='Resize', keep_ratio=True), dict(type='RandomFlip'), dict(type='Normalize', **img_norm_cfg), dict(type='Pad', size_divisor=32), dict(type='ImageToTensor', keys=['img']), dict(type='Collect', keys=['img'])])]
model = init_detector(config, CHECKPOINT_PATH, device=device)
model.eval()

from tqdm import tqdm
import os
results = []
with torch.no_grad():
    for img_name in tqdm(os.listdir(TEST_IMG_DIR)):
        img_pth = os.path.join(TEST_IMG_DIR, img_name)
        result = inference_detector(model, img_pth)
        boxes = result[0][:, :4]
        scores = result[0][:, 4]
        if len(boxes) > 0:
            boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
            boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
        result = {'image_id': img_name[:-4], 'PredictionString': format_prediction_string(boxes, scores)}
        results.append(result)

import pandas as pd
test_df = pd.DataFrame(results, columns=['image_id', 'PredictionString'])
test_df.to_csv('submission.csv', index=False)
test_df.head()