Introducción
Los sistemas de inspección de rayos X son fundamentales en el control de seguridad de aeropuertos, estaciones de tren y edificios públicos. Tradicionalmente, la identificación de artículos peligrosos depende de la inspección visual humana, lo que conlleva fatiga operadora y riesgo de error. Este artículo describe la implementación de un sistema de detección automática basado en YOLOv5 capaz de reconocer diez categorías de objetos prohibidos en imágenes de rayos X.
Conjunto de datos
Los artículos prohibidos se clasifican legalmente en múltiples grupos. Las categorías principales incluyen: armas de fuego, herramientas metálicas, animales vivos, productos de desinfección, sustancias corrosivas, combustibles y gases, materiales inflamables, alimentos perecederos, explosivos y objetos cortantes.
Para este proyecto se recopilaron 4.000 imágenes a color de rayos X provenientes de repositorios públicos y fuantes propias. Se etiquetaron diez categorías de objetos: encendedor, tijeras, batería portátil, recipiente a presión, cuchillo, combustible Zippo, esposas, honda, petardos y esmalte de uñas. La hreramienta LabelImg se utilizó con el formato de etiquetado de YOLO.
Configuración del entorno
conda create -n xray_detection python=3.8
conda activate xray_detection
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install opencv-python pyyaml pandas seaborn matplotlib tqdm pillow onnx tensorboard scipy
Configuración del modelo
Se debe modificar el archivo de arquitectura del modelo para ajustar el número de clases. A continuación se muestra una versión adaptada del archivo yolov5s.yaml:
# Parámetros de configuración
nc: 10 # número total de categorías
depth_multiple: 0.33
width_multiple: 0.50
# Anclas predefinidas
anchors:
- [10,13, 16,30, 33,23] # escala P3/8
- [30,61, 62,45, 59,119] # escala P4/16
- [116,90, 156,198, 373,326] # escala P5/32
# Red troncal
backbone:
[[-1, 1, Focus, [64, 3]],
[-1, 1, Conv, [128, 3, 2]],
[-1, 3, BottleneckCSP, [128]],
[-1, 1, Conv, [256, 3, 2]],
[-1, 9, BottleneckCSP, [256]],
[-1, 1, Conv, [512, 3, 2]],
[-1, 9, BottleneckCSP, [512]],
[-1, 1, Conv, [1024, 3, 2]],
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, BottleneckCSP, [1024, False]],
]
# Cabezal de detección
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]],
[-1, 3, BottleneckCSP, [512, False]],
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]],
[-1, 3, BottleneckCSP, [256, False]],
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]],
[-1, 3, BottleneckCSP, [512, False]],
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]],
[-1, 3, BottleneckCSP, [1024, False]],
[[17, 20, 23], 1, Detect, [nc, anchors]],
]
Archivo de configuración de datos
Crear el archivo prohibited_items.yaml en el directorio data/:
train: data/datasets/train.txt
val: data/datasets/val.txt
nc: 10
names: ['lighter','scissors','powerbank','pressure','knife','zippooil','handcuffs','slingshot','firecrackers','nailpolish']
Entrenamiento del modelo
Entrenamiento con una sola GPU:
python train.py --cfg models/yolov5s.yaml --data data/prohibited_items.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0
Entrenamiento con múltiples GPUs:
python train.py --cfg models/yolov5s.yaml --data data/prohibited_items.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0,1
Evaluación del modelo entrenado:
python val.py --weights runs/train/exp/weights/best.pt --data data/prohibited_items.yaml --device 0 --verbose
Interfaz de inferencia con PyQt5
Para el despliegue de la interfaz de inferencia se optó por PyQt5, que ofrece mayor flexibilidad que Gradio para construir interfaces complejas con múltiples controles.
pip install PyQt5
pip install PyQt5-tools
El siguiente código implementa una aplicación de detección con soporte para imagen, vídeo y cámara web:
import sys
import os
import numpy as np
import cv2
import torch
import torch.backends.cudnn as cudnn
from random import randint
from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import plot_one_box
from utils.torch_utils import select_device
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QFileDialog, QMessageBox
def resize_with_padding(img, target_dim=640, pad_color=(114, 114, 114)):
h_orig, w_orig = img.shape[:2]
scale_factor = min(target_dim / h_orig, target_dim / w_orig)
if scale_factor < 1.0:
scale_factor = min(scale_factor, 1.0)
new_w = int(round(w_orig * scale_factor))
new_h = int(round(h_orig * scale_factor))
resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
pad_w = (target_dim - new_w) / 2
pad_h = (target_dim - new_h) / 2
top = int(round(pad_h - 0.1))
bottom = int(round(pad_h + 0.1))
left = int(round(pad_w - 0.1))
right = int(round(pad_w + 0.1))
padded = cv2.copyMakeBorder(resized, top, bottom, left, right,
cv2.BORDER_CONSTANT, value=pad_color)
return padded, (scale_factor, scale_factor), (pad_w, pad_h)
class XRayDetectionApp(QtWidgets.QWidget):
def __init__(self, parent=None):
super(XRayDetectionApp, self).__init__(parent)
self.video_timer = QtCore.QTimer()
self.video_capture = cv2.VideoCapture()
self.camera_index = 0
self._setup_ui()
self._connect_signals()
self.conf_threshold = 0.5
self.iou_threshold = 0.45
self.input_size = 640
self.use_augment = False
self.compute_device = select_device('')
self.use_half = self.compute_device.type != 'cpu'
weight_path = os.path.join('weights', 'prohibited_best.pt')
self.detector = attempt_load(weight_path, device=self.compute_device)
self.input_size = check_img_size(self.input_size,
s=self.detector.stride.max())
if self.use_half:
self.detector.half()
cudnn.benchmark = True
self.class_names = (self.detector.module.names
if hasattr(self.detector, 'module')
else self.detector.names)
self.box_colors = [[randint(0, 255) for _ in range(3)]
for _ in range(len(self.class_names))]
def _setup_ui(self):
main_layout = QtWidgets.QHBoxLayout()
btn_layout = QtWidgets.QVBoxLayout()
self.btn_image = QtWidgets.QPushButton('Abrir imagen')
self.btn_camera = QtWidgets.QPushButton('Cámara')
self.btn_video = QtWidgets.QPushButton('Abrir video')
for btn in [self.btn_image, self.btn_camera, self.btn_video]:
btn.setMinimumHeight(50)
self.display_label = QtWidgets.QLabel()
self.display_label.setFixedSize(641, 481)
self.display_label.setAutoFillBackground(False)
btn_layout.addWidget(self.btn_image)
btn_layout.addWidget(self.btn_camera)
btn_layout.addWidget(self.btn_video)
main_layout.addLayout(btn_layout)
main_layout.addWidget(self.display_label)
self.setLayout(main_layout)
self.setWindowTitle('Detección de objetos en rayos X')
def _connect_signals(self):
self.btn_image.clicked.connect(self._on_image_clicked)
self.btn_camera.clicked.connect(self._on_camera_clicked)
self.btn_video.clicked.connect(self._on_video_clicked)
self.video_timer.timeout.connect(self._process_video_frame)
def _run_inference(self, raw_img):
display_img = raw_img.copy()
with torch.no_grad():
padded, _, _ = resize_with_padding(raw_img, self.input_size)
tensor_img = padded[:, :, ::-1].transpose(2, 0, 1)
tensor_img = np.ascontiguousarray(tensor_img)
tensor_img = torch.from_numpy(tensor_img).to(self.compute_device)
tensor_img = (tensor_img.half() if self.use_half
else tensor_img.float())
tensor_img /= 255.0
if tensor_img.ndimension() == 3:
tensor_img = tensor_img.unsqueeze(0)
predictions = self.detector(tensor_img, augment=self.use_augment)[0]
detections = non_max_suppression(
predictions, self.conf_threshold, self.iou_threshold)
for det in detections:
if det is not None and len(det):
det[:, :4] = scale_coords(
tensor_img.shape[2:], det[:, :4],
display_img.shape).round()
for *xyxy, conf, cls in reversed(det):
lbl = '%s %.2f' % (self.class_names[int(cls)], conf)
plot_one_box(xyxy, display_img, label=lbl,
color=self.box_colors[int(cls)],
line_thickness=3)
return display_img
def _display_result(self, img_result):
resized = cv2.resize(img_result, (640, 480),
interpolation=cv2.INTER_AREA)
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
qt_img = QtGui.QImage(rgb.data, rgb.shape[1], rgb.shape[0],
QtGui.QImage.Format_RGB888)
self.display_label.setPixmap(QtGui.QPixmap.fromImage(qt_img))
def _on_image_clicked(self):
path, _ = QFileDialog.getOpenFileName(
self, 'Seleccionar imagen', '',
'*.jpg;;*.png;;*.jpeg;;All Files(*)')
if not path:
return
img = cv2.imread(path)
if img is None:
return
result = self._run_inference(img)
self._display_result(result)
def _on_video_clicked(self):
if self.video_timer.isActive():
self.video_timer.stop()
self.video_capture.release()
self.display_label.clear()
self.btn_video.setText('Abrir video')
return
path, _ = QFileDialog.getOpenFileName(
self, 'Seleccionar video', '', '*.mp4;;*.avi;;All Files(*)')
if not path:
return
if not self.video_capture.open(path):
QMessageBox.warning(self, 'Aviso', 'No se pudo abrir el video')
return
self.btn_video.setText('Cerrar video')
self.video_timer.start(30)
def _on_camera_clicked(self):
if self.video_timer.isActive():
self.video_timer.stop()
self.video_capture.release()
self.display_label.clear()
self.btn_camera.setText('Cámara')
return
if not self.video_capture.open(self.camera_index):
QMessageBox.warning(self, 'Aviso',
'Verifique la conexión de la cámara')
return
self.btn_camera.setText('Cerrar cámara')
self.video_timer.start(30)
def _process_video_frame(self):
ret, frame = self.video_capture.read()
if frame is None:
self.video_timer.stop()
self.video_capture.release()
self.display_label.clear()
self.btn_video.setText('Abrir video')
self.btn_camera.setText('Cámara')
return
result = self._run_inference(frame)
self._display_result(result)
if __name__ == '__main__':
app = QtWidgets.QApplication(sys.argv)
window = XRayDetectionApp()
window.show()
sys.exit(app.exec_())
Consideraciones sobre precisión
El modelo YOLOv5s ofrece un equilibrio razonable entre velocidad y precisión para objetos pequeños como ancendedores en entornos densos. Para aplicaciones que requieran mayor exactitud, se puede recurrir a modelos más grandes (YOLOv5m o YOLOv5l) o migrar a YOLOv8. En contextos de seguridad, los falsos negativos son más críticos que los falsos positivos, por lo que conviene ajustar el umbral de confianza hacia valores más bajos y enriquecer el conjunto de datos con muestras negativas y objetos similares.