image_output_1.png
from PIL.ImageDraw import Draw

import PIL

import matplotlib.pyplot as plt

import numpy as np

import pandas as pd

import tensorflow as tf

import tensorflow_hub as hub

plt.style.use('fivethirtyeight')
plt.rcParams['figure.figsize'] = (12, 8)
plt.rcParams.update({'font.size': 14})
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.options.display.float_format = '{:.4f}'.format

META_FILE = '../input/crowd-counting/labels.csv'
data = pd.read_csv(META_FILE)

def reconstruct_path(image_id):
    """Превращает номерной ID изображения в относительный путь.
    """
    image_id = str(image_id).rjust(6, '0')
    return f'../input/crowd-counting/frames/frames/seq_{image_id}.jpg'
data['path'] = data['id'].apply(reconstruct_path)

MODEL_PATH = 'https://tfhub.dev/tensorflow/efficientdet/d0/1'
detector = hub.load(MODEL_PATH)

def detect_objects(path, model):
    """Извлекает изображение из файла, добавляет новую ось и применяет модель на объект.
    """
    image_tensor = tf.image.decode_jpeg(tf.io.read_file(path), channels=3)[tf.newaxis, ...]
    return model(image_tensor)
def count_persons(path, model, threshold=0.0):
    """Считает количество людей на изображении.
    """
    results = detect_objects(path, model)
    return (results['detection_classes'].numpy()[0] == 1)[np.where(results['detection_scores'].numpy()[0] > threshold)].sum()
def draw_boxes(image_path, data, threshold=0.0):
    """Возвращает изображения с прямоугольниками поверх каждого обнаруженного человека.
    """
    image = PIL.Image.open(image_path)
    draw = Draw(image)
    im_width, im_height = image.size
    boxes = data['detection_boxes'].numpy()[0]
    classes = data['detection_classes'].numpy()[0]
    scores = data['detection_scores'].numpy()[0]
    for i in range(int(data['num_detections'][0])):
        if classes[i] == 1 and scores[i] > threshold:
            ymin, xmin, ymax, xmax = boxes[i]
            left, right, top, bottom = (xmin * im_width, xmax * im_width, ymin * im_height, ymax * im_height)
            draw.line([(left, top), (left, bottom), (right, bottom), (right, top), (left, top)], width=4, fill='red')
    return image

example_path = '../input/crowd-counting/frames/frames/seq_000010.jpg'
results = detect_objects(example_path, detector)

draw_bboxes(example_path, results, threshold=0.25)