Machine Learning, Нейронные сети

Ищем «новый год» на картинках

Время прочтения: 3 мин.

Так как до нового года осталось чуть больше месяца, давайте попробуем сделать простую программу, которая позволит нам ответить на вопросы: как и когда Дед Мороз выполняет свои обязанности? Причем сделать это постараемся самым простым для нас способом.

Для реализации программы мы будем использовать предобученную нейросеть и простую классификацию.

Алгоритм работы программы будет следующим:

  1. Вначале загружаем картинку (кадр из видео потока);
  2. Применяем библиотеку ImageAI с предобученной моделью RetinaNet для получения классов найденных на изображении объектов;
  3. Обобщить название найденных классов для получения названия ситуации.

Разберем простой пример: предположим если мы видим диван, человека в красном костюме и дерево — то Дед мороз раздает подарки. Иначе главный персонаж наступающего праздника «еще отдыхает».

В начале установим все библиотеки для работы:

Tensorflow: pip install tensorflow
Numpy: pip install numpy
SciPy: pip install scipy
OpenCV: pip install opencv-python
Pillow: pip install pillow
Matplotlib: pip install matplotlib
H5py: pip install h5py
Keras: pip install keras
ImageAI: pip install imageai

Кроме этого нам необходимо скачать уже предобученную модель. Все варианты расположены на странице проекта ImageAI (ссылка). Мы скачаем библиотеку, которую предлагает автор тестового примера ImageAI под названием resnet50_coco_best_v2.0.1.h5.

После этого напишем простейший код для определения ситуации, в которой оказался Дед Мороз.

Для этого мы сначала импортируем необходимые для нас модули:

from imageai.Detection import ObjectDetection
import os

После этого укажем в переменной папку, в которой мы работаем:

execution_path = os.getcwd()

и создадим объект-детектор с загруженной ранее моделью:

detector = ObjectDetection()
detector.setModelTypeAsRetinaNet()
detector.setModelPath( os.path.join(execution_path , "resnet50_coco_best_v2.0.1.h5"))
detector.loadModel()

Последним этапом работы с изображением будет загрузка его в детектор и создание нового изображения, на котором будут отмечены все найденные объекты:

detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path , 
         "image.jpg"), output_image_path=os.path.join(execution_path , "imagenew.jpg"))

Названия классов, найденных на изображении, будут находится в переменной eachObject. Для проверки, в консоли отобразим полный список найденного.

all_object  = []
for eachObject in detections:
    all_object.append(eachObject["name"])
    print(eachObject["name"] + " : " + eachObject["percentage_probability"] )

Но мы помним, что нашей задачей является не только определение объектов, но и определение ситуации. Для этого пройдемся еще раз по всем найденным объектам и проверим есть ли классы, которые мы хотим найти:

class_new_year = ["person","tree","furniture"]
result=list(set(array_new_year) & set(all_object))
if len(result) == 3:
	print("Дед Мороз вручает подарки!")
else:
	print("Дед Мороз еще не дарит подарки((( ")

В итоге загрузим картинку:

Результат работы нашего алгоритма будет следующий:

person : 97.45132427196611
Дед Мороз еще не дарит подарки(((

На данной картинке как мы видим главный герой наступающего праздника не дарит подарки. Наша программа правильно определила ситуацию. При этом данный алгоритм не требует больших ресурсов от компьютера, и позволяет обрабатывать кадры из видео с небольшой задержкой. 

Вот таким простым образом, можно быстро определить ситуацию, которая происходит на изображении. При этом для расширения списка ситуаций мы можем как увеличить количество массивов, так и использовать более сложные методы Text Mining.

Советуем почитать