Improved Co-DETR With Dropkey and Its Application to Hot Work Detection

Abstract

Although ViT has achieved significant success in the field of image classification, research on ViT-based object detection algorithms is still in its early stages, and their application in real-world scenarios is limited. Furthermore, algorithms based on ViT or Transformer are prone to overfitting issues when training data is scarce. While CO-DETR has achieved state-of-the-art object detection precision on the COCO dataset leaderboard, the ViT-based CO-DETR also suffers from overfitting problems, which affect its detection precision on smaller datasets. Based on the study of ViT-based object detection algorithms, a new object detection algorithm termed DC-DETR (DropKey Co-DETR) was proposed in this paper. It builds upon CO-DETR and introduces a regularization method called DropKey into the Transformer attention mechanism. By randomly dropping part of the Key during the attention phase, the network is encouraged to capture global information about the target object. This method effectively alleviates the overfitting problem in ViT for object detection tasks, improving the model's precision and generalization ability. To validate the effectiveness and practical applicability of DC-DETR in environments with limited computational resources, a dataset for hot work scenarios was collected and annotated. Based on this dataset, performance tests were conducted on the DC-DETR, CO-DETR, and YOLOv5 algorithms. The test results indicate that the proposed DC-DETR algorithm exhibits superior performance, with detection precision improving by 0.7% compared to CO-DETR and by 5.7% compared to YOLOv5. The detection speed is the same as CO-DETR, and only 2.9 ms slower than YOLOv5. The experiments demonstrate that the proposed DC-DETR algorithm achieves a balance between precision and speed, making it well-suited for practical object detection applications.