Reconocimiento de objetos con TensorFlow y OpenCV

Xukyo

hace 10 meses

En este tutorial veremos cómo realizar reconocimiento de objetos con TensorFlow y OpenCV usando una red neuronal pre-entrenada usando deep learning.

En un tutorial anterior vimos cómo reconocer formas simples mediante visión por ordenador. Este método sólo funciona para ciertas formas simples predefinidas. Si quieres reconocer una mayor variedad de objetos, la forma más sencilla es utilizar la inteligencia artificial.

Hardware

Un ordenador con una instalación de Python3
Una cámara

Principio

La inteligencia artificial es un campo de la informática en el que el propio programa aprende a realizar determinadas tareas. En particular, el reconocimiento visual. En este tutorial, vamos a utilizar una red neuronal entrenada para reconocer determinadas formas.

Se necesitan muchos datos para poder entrenar correctamente una red neuronal. Se ha demostrado que el aprendizaje es más rápido en una red neuronal entrenada para otra cosa. Por ejemplo, una red neuronal entrenada para reconocer perros se entrenará más fácilmente para reconocer gatos.

Configuración de Python

Si no, puedes descargar e instalar Python 3

A continuación, puede instalar las bibliotecas OpenCV, numpy e imutils necesarias.

python3 -m pip install opencv-python numpy imutils

Éstas son las versiones que utilizo en este tutorial

numpy== 1.22.0
tensorflow== 2.13.0
opencv== 4.6.0

N.B.: Al utilizar ciertos paquetes python, como TensorFlow, pueden surgir con mucha frecuencia problemas de compatibilidad. Si tienes dificultades, no dudes en instalar versiones específicas de los paquetes. Si tienes varios proyectos en curso, te recomiendo encarecidamente que crees entornos virtuales (venv).

Recuperar un modelo preentrenado

Descarga del modelo ModelNet-SSD del Model Zoo (por ejemplo, MobileNet SSD v2 320×320)

Descarga el fichero mscoco_label_map.pbtxt, que contiene los identificadores de los objetos reconocidos por el modelo.

N.B.: en un tutorial anterior, utilizamos los resultados del framework Caffe para el modelo (archivos prototxt y caffemodel). En este artículo, utilizamos SavedModel de TensorFlow.

Crea una carpeta TensorFlowCV, que será tu espacio de trabajo. En esta carpeta, crea una carpeta pretrained_models en la que podrás descomprimir el modelo descargado.

También puedes crear una carpeta de datos en la que colocar tus imágenes o vídeos.

Coloca los archivos del modelo en una carpeta y crea el archivo ObjectRecognition.py

Script en Python para el reconocimiento de objetos

En primer lugar, creamos un flujo de vídeo (vs) utilizando la biblioteca imutils, que recuperará las imágenes de la cámara.

vs = VideoStream(src=0, resolution=(1600, 1200)).start()

Inicializamos una red neuronal con los parámetros de SSD-ModelNetV2 (net) utilizando la biblioteca TensorFlow.

model= tf.saved_model.load("./pretrained_models/ssd_mobilenet_v2_320x320_coco17_tpu-8/saved_model")

Una función recupera del fichero pbtxt los nombres de las clases reconocidas por el modelo

#load class names
def read_label_map(label_map_path):

    item_id = None
    item_name = None
    items = {}
    
    with open(label_map_path, "r") as file:
        for line in file:
            line.replace(" ", "")
            if line == "item{":
                pass
            elif line == "}":
                pass
            elif "id" in line:
                item_id = int(line.split(":", 1)[1].strip())
            elif "display_name" in line: #elif "name" in line:
                item_name = line.split(":", 1)[1].replace("'", "").strip()

            if item_id is not None and item_name is not None:
                #items[item_name] = item_id
                items[item_id] = item_name
                item_id = None
                item_name = None

    return items

class_names=read_label_map("./pretrained_models/ssd_mobilenet_v2_320x320_coco17_tpu-8/mscoco_label_map.pbtxt")
class_colors = np.random.uniform(0, 255, size=(len(class_names), 3))

A continuación, crearemos un bucle que, en cada iteración, leerá la imagen de la cámara y la pasará a la entrada de la red neuronal para la detección y el reconocimiento de objetos.

	#Main loop
	while True:
		# Get video sttream. max width 800 pixels 
		#img = vs.read()
		img= cv2.imread('./data/two-boats.jpg') #from image file
		#ret, img=vc.read() #from video or ip cam
		img = imutils.resize(img, width=800)

		img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
		# get height and width of image
		h, w, _ = img.shape

		input_tensor = np.expand_dims(img, 0)

		# predict from model
		resp = model(input_tensor)

Por último, el código muestra el cuadro de detección en la imagen, la probabilidad de reconocimiento y la posición.

					# write classname for bounding box
					cls=int(cls) #convert tensor to index
					label = "{}: {:.2f}%".format(class_names[cls],score * 100)
					cv2.putText(img, label, (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 1, class_colors[cls], 1)
					
					#display position
					X= (xmax+xmin)/2
					Y= (ymax+ymin)/2
					poslbl= "X: ({},{})".format(X,Y)
					cv2.circle(img, (int(X)-15, int(Y)), 1, class_colors[cls], 2)	
					cv2.putText(img, poslbl, (int(X), int(Y)),
						cv2.FONT_HERSHEY_SIMPLEX, 0.5, class_colors[cls], 2)
					
					# draw on image
					cv2.rectangle(img, (xmin, ymin), (xmax, ymax), class_colors[cls], 4)

Código completo de reconocimiento de objetos con OpenCV y TensorFlow

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  ObjectRecognitionTFVideo.py
#  Description:
#		Use ModelNetV2-SSD model to detect objects on image or video
#
#  www.aranacorp.com

# import packages
import sys
from imutils.video import VideoStream
from imutils.video import FPS
import numpy as np
import argparse
import imutils
import time
import cv2
import tensorflow as tf


# load model from path
model= tf.saved_model.load("./pretrained_models/ssd_mobilenet_v2_320x320_coco17_tpu-8/saved_model")
print("model loaded")

#load class names
def read_label_map(label_map_path):

    item_id = None
    item_name = None
    items = {}
    
    with open(label_map_path, "r") as file:
        for line in file:
            line.replace(" ", "")
            if line == "item{":
                pass
            elif line == "}":
                pass
            elif "id" in line:
                item_id = int(line.split(":", 1)[1].strip())
            elif "display_name" in line: #elif "name" in line:
                item_name = line.split(":", 1)[1].replace("'", "").strip()

            if item_id is not None and item_name is not None:
                #items[item_name] = item_id
                items[item_id] = item_name
                item_id = None
                item_name = None

    return items

class_names=read_label_map("./pretrained_models/ssd_mobilenet_v2_320x320_coco17_tpu-8/mscoco_label_map.pbtxt")
class_colors = np.random.uniform(0, 255, size=(len(class_names), 3))


if __name__ == '__main__':

	# Camera initialisation
	print("Start Camera...")
	vs = VideoStream(src=0, resolution=(1600, 1200)).start() #from usb cam
	#vs = VideoStream(usePiCamera=True, resolution=(1600, 1200)).start() #from RPi cam
	#vc = cv2.VideoCapture('./data/Splash - 23011.mp4') #from video file

	time.sleep(2.0)
	fps = FPS().start()
	
	#Main loop
	while True:
		#get image
		img = vs.read() # Get video stream
		#img= cv2.imread('./data/two-boats.jpg') #from image file
		#ret, img=vc.read() #from video or ip cam

                #process image
		img = imutils.resize(img, width=800) #max width 800 pixels 
		img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
		# get height and width of image
		h, w, _ = img.shape

		input_tensor = np.expand_dims(img, 0)

		# predict from model
		resp = model(input_tensor)

		# iterate over boxes, class_index and score list
		for boxes, classes, scores in zip(resp['detection_boxes'].numpy(), resp['detection_classes'], resp['detection_scores'].numpy()):
			for box, cls, score in zip(boxes, classes, scores): # iterate over sub values in list
				if score &gt; 0.61: # we are using only detection with confidence of over 0.6
					ymin = int(box[0] * h)
					xmin = int(box[1] * w)
					ymax = int(box[2] * h)
					xmax = int(box[3] * w)
									
					# write classname for bounding box
					cls=int(cls) #convert tensor to index
					label = "{}: {:.2f}%".format(class_names[cls],score * 100)
					cv2.putText(img, label, (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 1, class_colors[cls], 1)
					
					#display position
					X= (xmax+xmin)/2
					Y= (ymax+ymin)/2
					poslbl= "X: ({},{})".format(X,Y)
					cv2.circle(img, (int(X)-15, int(Y)), 1, class_colors[cls], 2)	
					cv2.putText(img, poslbl, (int(X), int(Y)),
						cv2.FONT_HERSHEY_SIMPLEX, 0.5, class_colors[cls], 2)
					
					# draw on image
					cv2.rectangle(img, (xmin, ymin), (xmax, ymax), class_colors[cls], 4)
				
				
		# Show video frame
		cv2.imshow("Frame", img)
		key = cv2.waitKey(1) &amp; 0xFF

		# Exit script with letter q
		if key == ord("q"):
			break

		# FPS update 
		fps.update()

	# Stops fps and display info
	fps.stop()
	print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
	print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))

	cv2.destroyAllWindows()
	vs.stop()
	#vc.release()

Fuentes de imágenes para la detección de objetos

Puede utilizar este script con diferentes fuentes de imágenes. Para ello, debe adaptar ligeramente el código anterior para modificar la variable «img» que contiene la imagen que se va a analizar.

La cámara web de tu ordenador

vs = VideoStream(src=0, resolution=(1600, 1200)).start()
while True:
	frame = vs.read()

El flujo de vídeo debe detenerse al final del script con vs.stop()

Una cámara IP

vc = cv2.VideoCapture('rtsp://user:password@ipaddress:rtspPort')
while True:
	ret, frame=vc.read() #from ip cam

Asegúrate de detener la captura de vídeo al final del script con vc.release()

La Raspberry Pi Picam

vs = VideoStream(usePiCamera=True, resolution=(1600, 1200)).start()
while True:
	frame = vs.read()

Recuerda detener el flujo al final del script con vs.stop()

Un archivo de vídeo

vc = cv2.VideoCapture('./img/Splash - 23011.mp4') #from video
while True:
	ret, frame=vc.read() #from video

Un archivo de imagen

frame= cv2.imread('./img/two-boats.jpg')

Resultados

Para este ejemplo, enviamos una imagen de dos barcos como entrada a la red neuronal, que son reconocidos correctamente. Para obtener resultados ligeramente diferentes, puede modificar el parámetro de confianza para evitar falsos positivos.

Puedes probar este código con tu webcam o con fotos, por ejemplo, para ver cómo funcionan el modelo y el reconocimiento de objetos.

Una vez que tu script funcione, puedes entrenar tu modelo para detectar otros objetos.

Paquetes y modelos

En este tutorial, hemos utilizado el modelo preentrenado SSD ModelNetV2. Cabe destacar que existen otros modelos de clasificación con diferentes prestaciones y características.

vgg16
vgg19
resnet50
resnet101
resnet152
densenet121
densenet169
densenet201
inceptionresnetv2
inceptionv3
mobilenet
mobilenetv2
nasnetlarge
nasnetmóvil
xcepción

No dudes en dejarnos un comentario para compartir los modelos que utilizas o conoces y tu opinión.