inference demo

원본 : demo/image_demo.py

# Copyright (c) OpenMMLab. All rights reserved.
import asyncio
from argparse import ArgumentParser

from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('img', help='Image file')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument('--out-file', default=None, help='Path to output file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--palette',
        default='coco',
        choices=['coco', 'voc', 'citys', 'random'],
        help='Color palette used for visualization')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='bbox score threshold')
    parser.add_argument(
        '--async-test',
        action='store_true',
        help='whether to set async options for async inference.')
    args = parser.parse_args()
    return args


def main(args):
    # build the model from a config file and a checkpoint file
    model = init_detector(args.config, args.checkpoint, device=args.device)
    # test a single image
    result = inference_detector(model, args.img)
    # show the results
    show_result_pyplot(
        model,
        args.img,
        result,
        palette=args.palette,
        score_thr=args.score_thr,
        out_file=args.out_file)


async def async_main(args):
    # build the model from a config file and a checkpoint file
    model = init_detector(args.config, args.checkpoint, device=args.device)
    # test a single image
    tasks = asyncio.create_task(async_inference_detector(model, args.img))
    result = await asyncio.gather(tasks)
    # show the results
    show_result_pyplot(
        model,
        args.img,
        result[0],
        palette=args.palette,
        score_thr=args.score_thr,
        out_file=args.out_file)


if __name__ == '__main__':
    args = parse_args()
    if args.async_test:
        asyncio.run(async_main(args))
    else:
        main(args)

수정한 demo_img.py 참고

import warnings
from pathlib import Path
import numpy as np
import torch

import mmcv
from mmcv.ops      import RoIPool
from mmcv.parallel import collate, scatter
from mmcv.runner   import load_checkpoint
from mmcv.runner   import load_checkpoint

from mmdet.core               import get_classes
from mmdet.datasets           import replace_ImageToTensor
from mmdet.datasets.pipelines import Compose
from mmdet.models             import build_detector

from mmdet.apis import init_detector
from mmdet.apis import inference_detector
from mmdet.apis import show_result_pyplot

import os
INPUT_IMG = "./imgs/demo.jpg"
OUTPUT_FILE = "./imgs/demo_out.jpg"

CONFIG = [os.path.join('./param', filenm) for filenm in os.listdir('./param') if filenm.endswith('.py')][0]
CHKPNT = [os.path.join('./param', filenm) for filenm in os.listdir('./param') if filenm.endswith('.pth')][0] 
# ROOT = "/home/oschung_skcc/my/git/mmdetection/"
# CONFIG      = os.path.join(ROOT, "sixxconfigs/faster_rcnn_r50_fpn_1x_coco_001.py")
# CHECKPOINT  = os.path.join(ROOT, "checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth")

SCORE_THRESHOLD = 0.9
DEVICE  = "cuda:6"
PALETTE = "coco" # choices=['coco', 'voc', 'citys', 'random']

# def parse_args():
#     parser = ArgumentParser()
#     parser.add_argument('img',                          help='Image file')
#     parser.add_argument('config',                       help='Config file')
#     parser.add_argument('checkpoint',                   help='Checkpoint file')
#     parser.add_argument('--out-file', default=None,     help='Path to output file')
#     parser.add_argument('--device',   default='cuda:0', help='Device used for inference')
#     parser.add_argument('--palette',  default='coco',  choices=['coco', 'voc', 'citys', 'random'], help='Color palette used for visualization')
#     parser.add_argument('--score-thr', type=float, default=0.3, help='bbox score threshold')
#     parser.add_argument('--async-test', action='store_true',    help='whether to set async options for async inference.')
#     args = parser.parse_args()
#     return args

if __name__ == '__main__':
    # build the MODEL from a config file & a checkpoint file
    model = init_detector(CONFIG, CHKPNT, device=DEVICE)
    # test a single image
    result = inference_detector(model, INPUT_IMG)
    # show the results
    show_result_pyplot(model, INPUT_IMG, result,
        palette   = PALETTE,
        score_thr = SCORE_THRESHOLD,
    )
    # model.show_result(INPUT_IMG, result,
    #     score_thr=SCORE_THRESHOLD,
    #     show=True,
    #     #wait_time=wait_time,
    #     #win_name=title,
    #     bbox_color=PALETTE,
    #     text_color=(200, 200, 100),
    #     mask_color=PALETTE,
    #     out_file=OUTPUT_FILE)

import json
# !w g e t http://images.cocodataset.org/annotations/annotations_trainval2017.zip
# !unzip annotations_trainval2017.zip
with open('./annotations/instances_train2017.json', 'r') as f:
    data = json.load(f)
categories = data['categories']
category_list = []
for category in categories:
    category_list.append({'name': category['name'], 'id': category['id']})
# print(category_list)
classes = category_list

# classes = [
#     { "id": 1, "name": "person" },
#     { "id": 2, "name": "bicycle" },
#     { "id": 3, "name": "car" },
# ...
#     { "id":88, "name": "teddy_bear" },
#     { "id":89, "name": "hair_drier" },
#     { "id":90, "name": "toothbrush" }
# ]

encoded_results = []
idx = 0
for dets in result:
    if len(dets) == 0:
        idx = idx + 1
        continue
    for det in dets:
        if det[4] < SCORE_THRESHOLD:
            continue

        x_bbox = float(det[0])
        y_bbox = float(det[1])
        width  = float(det[2])
        height = float(det[3])
        score = float(det[4])
        
        encoded_results.append({
            'confidence': score,
            'label': classes[idx]['name'],
            'points': [x_bbox, y_bbox, width, height],
            'type': 'rectangle'
        })
    idx = idx + 1
    # if idx == 1:
    #     break

#-----------------------------------------------------------------
img = mmcv.imread(INPUT_IMG)

if isinstance(result, tuple):
    bbox_result, segm_result = result
    if isinstance(segm_result, tuple):
        segm_result = segm_result[0]  # discard the `dim`
else:
    bbox_result, segm_result = result, None

img = img.copy()
bboxes = np.vstack(bbox_result)
labels = [
    np.full(bbox.shape[0], i, dtype=np.int32)
    for i, bbox in enumerate(bbox_result)
]
labels = np.concatenate(labels)
scores = bboxes[:, -1]

inds = scores > SCORE_THRESHOLD
bboxes = bboxes[inds, :] # points
labels = labels[inds]    # label
scores = scores[inds]    # confidence

encoded_results = []
if bboxes.shape[0] > 0:
    for i in range(bboxes.shape[0]):
        encoded_results.append({
            'confidence': scores[i],
            'label': classes[labels[i]]['name'],
            'points': bboxes[i].tolist(),
            'type': 'rectangle'
        })

print(encoded_results)
# if show_mask and segm_result is not None:
#     segms = mmcv.concat_list(segm_result)
#     segms = [segms[i] for i in np.where(inds)[0]]
#     #colors = [next(palette) for _ in range(len(segms))]

model.show_result in show_result_pyplot

source code for model.show_result() in MMDetection

def show_result(img,
                result,
                palette=None,
                score_thr=0.3,
                show=True,
                out_file=None,
                thickness=1,
                font_scale=0.5,
                win_name='',
                show_mask=True):
    """Visualize the detection results on the image.

    Args:
        img (str or ndarray): The image to be displayed.
        result (tuple[list] or list): The detection result, can be either
            a tuple containing bbox and segm results or a list of bbox result.
        palette (list[list[int]]]: The color palette of masks.
        score_thr (float): The threshold to visualize the bboxes and masks.
        show (bool): Whether to show the image, default True.
        out_file (str or None): The filename to write the image.
        thickness (int): Thickness of lines.
        font_scale (float): Font scales of texts.
        win_name (str): The window name.
        show_mask (bool): Whether to display masks.

    Returns:
        ndarray: The image with bboxes and masks drawn on it.
    """
    img = mmcv.imread(img)
    if isinstance(result, tuple):
        bbox_result, segm_result = result
        if isinstance(segm_result, tuple):
            segm_result = segm_result[0]  # discard the `dim`
    else:
        bbox_result, segm_result = result, None
    img = img.copy()
    bboxes = np.vstack(bbox_result)
    labels = [
        np.full(bbox.shape[0], i, dtype=np.int32)
        for i, bbox in enumerate(bbox_result)
    ]
    labels = np.concatenate(labels)
    scores = bboxes[:, -1]
    inds = scores > score_thr
    bboxes = bboxes[inds, :]
    labels = labels[inds]
    scores = scores[inds]
    if show_mask and segm_result is not None:
        segms = mmcv.concat_list(segm_result)
        segms = [segms[i] for i in np.where(inds)[0]]
        if palette is None:
            palette = color_val_iter()
        colors = [next(palette) for _ in range(len(segms))]
    if bboxes.shape[0] > 0:
        if hasattr(palette, '__iter__'):
            color = next(palette)
        else:
            color = palette
        img = imshow_det_bboxes(
            img,
            bboxes,
            labels,
            scores,
            color=color,
            thickness=thickness,
            font_scale=font_scale,
            show=False)
        if show_mask and segm_result is not None:
            img = imshow_det_masks(
                img, segms, colors, alpha=0.5, show=False)
    if show:
        imshow(img, win_name, out_file=out_file)
    return img

main.py ver1.0

# Copyright Bong-Kyo Seo. All rights reserved.
from mmdet.apis import (inference_detector,
                        init_detector)
import json
import base64
from PIL import Image
import io
import numpy as np

config = './fifa.py'
checkpoint = './fifa.pth'
classes =  [
        { "id": 0, "name": "player" }
      ]

def init_context(context):
    context.logger.info("Init context... 0%")
    model = init_detector(config, checkpoint, device='cpu')
    context.user_data.model = model
    context.logger.info("Init context... 100%")

def handler(context, event):
    context.logger.info("Run yolox model")
    data = event.body
    buf = io.BytesIO(base64.b64decode(data["image"]))
    threshold = float(data.get("threshold", 0.9))
    context.user_data.model.conf = threshold
    image = Image.open(buf)
    imgArray = np.array(image)
    result = inference_detector(context.user_data.model, imgArray)

    encoded_results = []
    idx = 0
    for dets in result:
        if len(dets) == 0:
            idx = idx + 1
            continue
        for det in dets:
            if det[4] < 0.05:
                continue
                
            encoded_results.append({
                'confidence': float(det[4]),
                'label': classes[idx]['name'],
                'points': [float(det[0]), float(det[1]), float(det[2]), float(det[3])],
                'type': 'rectangle'
            })
        idx = idx + 1
        if idx == 1:
            break

    return context.Response(
        body=json.dumps(encoded_results),
        headers={},
        content_type='application/json',
        status_code=200
    )

main.py ver.2.0

# Copyright onesixx. All rights reserved.
import os 
CONFIG = [os.path.join('./param', filenm) for filenm in os.listdir('./param') if filenm.endswith('.py')][0]
CHKPNT = [os.path.join('./param', filenm) for filenm in os.listdir('./param') if filenm.endswith('.pth')][0] 
SCORE_THRESHOLD = 0.9
DEVICE  = "cpu" #"cuda:6"
PALETTE = "coco"

import json
with open('./annotations/instances_train2017.json', 'r') as f:
    data = json.load(f)
categories = data['categories']

category_list = []
for category in categories:
    category_list.append({'id':category['id'], 'name':category['name']})

classes = category_list
# classes = [
#     { "id": 1, "name": "person" },
#     { "id": 2, "name": "bicycle" },
# ...
#     { "id":90, "name": "toothbrush" }
# ]
import mmcv
from mmdet.apis import init_detector
from mmdet.apis import inference_detector

import io
import base64
from PIL import Image
import numpy as np

def init_context(context):
    context.logger.info("Init context... 0%") # --------------------------------
    model = init_detector(CONFIG, CHECKPOINT, device=DEVICE)
    context.user_data.model = model
    context.logger.info("Init context... 100%") # ------------------------------

def handler(context, event):
    context.logger.info("Run sixx model")
    data = event.body
    buf = io.BytesIO(base64.b64decode(data["image"]))

    # threshold = float(data.get("threshold", SCORE_THRESHOLD))
    # context.user_data.model.conf = threshold

    image = Image.open(buf)
    imgArray = np.array(image)

    result = inference_detector(context.user_data.model, imgArray)

    if isinstance(result, tuple):
        bbox_result, segm_result = result
        if isinstance(segm_result, tuple):
            segm_result = segm_result[0]  # discard the `dim`
    else:
        bbox_result, segm_result = result, None

    img = image.copy()
    bboxes = np.vstack(bbox_result)
    labels = [
        np.full(bbox.shape[0], i, dtype=np.int32)
        for i, bbox in enumerate(bbox_result)
    ]
    labels = np.concatenate(labels)
    scores = bboxes[:, -1]

    inds = scores > SCORE_THRESHOLD
    scores = scores[inds]    # confidence
    labels = labels[inds]    # label
    bboxes = bboxes[inds, :] # points

    if show_mask and segm_result is not None:
        segms = mmcv.concat_list(segm_result)
        segms = [segms[i] for i in np.where(inds)[0]]
        if palette is None:
            palette = color_val_iter()
        colors = [next(palette) for _ in range(len(segms))]

    encoded_results = []
    if bboxes.shape[0] > 0:
        for i in range(bboxes.shape[0]):
            encoded_results.append({
                'confidence': scores[i],
                'label': classes[labels[i]]['name'],
                'points': bboxes[i].tolist(),
                'type': 'rectangle'
            })
    return context.Response(
        body=json.dumps(encoded_results),
        headers={},
        content_type='application/json',
        status_code=200
    )