【物体検出】YOLOXまとめ｜第6回：物体検出を応用してモザイク処理を自動化してみる

from google.colab import drive
drive.mount('/content/drive')
%cd ./drive/MyDrive

!git clone https://github.com/Megvii-BaseDetection/YOLOX

%cd YOLOX
!pip install -U pip && pip install -r requirements.txt
!pip install -v -e . 
!pip install cython
!pip install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth

!python tools/export_onnx.py --output-name yolox_x.onnx -n yolox-x -c yolox_x.pth

import argparse
import os
import cv2
import numpy as np
import pandas as pd
import onnxruntime
from google.colab.patches import cv2_imshow
from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis

output_dir ='onnx_out'
image_path = 'sample.jpg'
model = 'yolox_x.onnx'

input_shape = (640,640)
origin_img = cv2.imread(image_path)
img, ratio = preprocess(origin_img, input_shape)
session = onnxruntime.InferenceSession(model)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.5)
if dets is not None:
    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
    origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                      0.3, class_names=COCO_CLASSES)
mkdir(output_dir)
output_path = os.path.join(output_dir, os.path.basename(image_path))

cv2.imwrite(output_path, origin_img)
cv2_imshow(origin_img)

result = []
[result.extend((final_cls_inds[x],COCO_CLASSES[int(final_cls_inds[x])],final_scores[x],final_boxes[x][0],final_boxes[x][1],final_boxes[x][2],final_boxes[x][3]) for x in range(len(final_scores)))]
df = pd.DataFrame(result, columns = ['class-id','class','score','x-min','y-min','x-max','y-max'])
df

def mosaic(img, alpha):
    w = img.shape[1]
    h = img.shape[0]
    img = cv2.resize(img, (int(w*alpha), int(h*alpha)))
    img = cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)
    return img

input_shape = (640,640)
origin_img = cv2.imread(image_path)
img, ratio = preprocess(origin_img, input_shape)
session = onnxruntime.InferenceSession(model)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.5)

mkdir(output_dir)
output_path = os.path.join(output_dir, os.path.basename(image_path))
final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]

for i in range(len(final_scores)):
  xmin = int(final_boxes[i][0])
  ymin = int(final_boxes[i][1])
  xmax = int(final_boxes[i][2])
  ymax = int(final_boxes[i][3])
  origin_img[ymin:ymax,xmin:xmax] = mosaic(origin_img[ymin:ymax,xmin:xmax],0.1)
  print(xmin,ymin,xmax,ymax,final_cls_inds[i])

origin_img[ymin:ymax,xmin:xmax] = mosaic(origin_img[ymin:ymax,xmin:xmax],0.1)