python - Yolov3.weights 未完全读取错误

问题描述

从 main.py 文件运行代码时，在未读取 yolov3.weights 文件的 utils.py 文件中出现错误。此错误出现在 load_yolo_weights 函数中的 utils.py 文件中。weights_file 将程序定向到正确的文件位置。YOLO_TYPE 是 yolov3

下面提供的是代码。

主要.py：

import cv2
import numpy as np
import argparse
import sys
import os
import tensorflow as tf
from yolov3.utils import Load_Yolo_model

# starting deep sort
maxCosineDist = 0.7
nn_budget = None

modelFilename = 'mars-small128.pb'
encoder = gendetect.create_box_encoder(modelFilename, batch_size=1)
metric = nn_matching.NearestNeighborDistanceMetric("cosine", maxCosineDist, nn_budget)
tracker = Tracker(metric)

yolo = Load_Yolo_model()

# Read image

out = "Yolo_Output.avi"
if args.image:
    if not os.path.isfile(args.image):
        print(args.image, " doesn't exist")
        sys.exit(1)
    vid = cv2.VideoCapture(args.image)
    out = args.image[:-4] + "Yolo_Output.jpg"

# Load class names
classesFile = "coco.names"
classes = read_class_names(classesFile)

keyList = list(classes.keys())
valList = list(classes.values())

# Create window for output
windowTitle = "Sports Tracking and Analytics"
cv2.namedWindow(windowTitle, cv2.WINDOW_NORMAL)


while True:
    _, frame = vid.read()

    try:
        origFrame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        origFrame = cv2.cvtColor(origFrame, cv2.COLOR_BGR2RGB)

    except:
        break

    imgData = image_preprocess(np.copy(origFrame), [416, 416])
    imgData = imgData[np.newaxis, ...].astype(np.float32)

    predBox = yolo.predict(imgData)

实用程序.py

import cv2
import time
import random
import colorsys

from yolov3.yolov4 import *
from yolov3.configs import *
import tensorflow as tf
import numpy as np
from tensorflow.python.saved_model import tag_constants

def load_yolo_weights(model, weights_file):
    tf.keras.backend.clear_session() # used to reset layer names
    # load Darknet original weights to TensorFlow model
    if YOLO_TYPE == "yolov3":
        range1 = 75 if not TRAIN_YOLO_TINY else 13
        range2 = [58, 66, 74] if not TRAIN_YOLO_TINY else [9, 12]
    if YOLO_TYPE == "yolov4":
        range1 = 110 if not TRAIN_YOLO_TINY else 21
        range2 = [93, 101, 109] if not TRAIN_YOLO_TINY else [17, 20]

    with open(weights_file, 'rb') as wf:
        major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)

        j = 0
        for i in range(range1):
            if i > 0:
                conv_layer_name = 'conv2d_%d' %i
            else:
                conv_layer_name = 'conv2d'

            if j > 0:
                bn_layer_name = 'batch_normalization_%d' %j
            else:
                bn_layer_name = 'batch_normalization'

            conv_layer = model.get_layer(conv_layer_name)
            filters = conv_layer.filters
            k_size = conv_layer.kernel_size[0]
            in_dim = conv_layer.input_shape[-1]

            if i not in range2:
                # darknet weights: [beta, gamma, mean, variance]
                bn_weights = np.fromfile(wf, dtype=np.float32, count=4 * filters)
                # tf weights: [gamma, beta, mean, variance]
                bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]
                bn_layer = model.get_layer(bn_layer_name)
                j += 1
            else:
                conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)

            # darknet shape (out_dim, in_dim, height, width)
            conv_shape = (filters, in_dim, k_size, k_size)
            conv_weights = np.fromfile(wf, dtype=np.float32, count=np.product(conv_shape))
            # tf shape (height, width, in_dim, out_dim)
            conv_weights = conv_weights.reshape(conv_shape).transpose([2, 3, 1, 0])

            if i not in range2:
                conv_layer.set_weights([conv_weights])
                bn_layer.set_weights(bn_weights)
            else:
                conv_layer.set_weights([conv_weights, conv_bias])

        assert len(wf.read()) == 0, 'failed to read all data'

def Load_Yolo_model():
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if len(gpus) > 0:
        print(f'GPUs {gpus}')
        try: tf.config.experimental.set_memory_growth(gpus[0], True)
        except RuntimeError: pass

    if YOLO_FRAMEWORK == "tf": # TensorFlow detection
        if YOLO_TYPE == "yolov4":
            Darknet_weights = YOLO_V4_TINY_WEIGHTS if TRAIN_YOLO_TINY else YOLO_V4_WEIGHTS
        if YOLO_TYPE == "yolov3":
            Darknet_weights = YOLO_V3_TINY_WEIGHTS if TRAIN_YOLO_TINY else YOLO_V3_WEIGHTS

        if YOLO_CUSTOM_WEIGHTS == False:
            print("Loading Darknet_weights from:", Darknet_weights)
            yolo = Create_Yolo(input_size=416, CLASSES="coco.names")
            load_yolo_weights(yolo, Darknet_weights) # use Darknet weights
        else:
            checkpoint = f"./checkpoints/{TRAIN_MODEL_NAME}"
            if TRAIN_YOLO_TINY:
                checkpoint += "_Tiny"
            print("Loading custom weights from:", checkpoint)
            yolo = Create_Yolo(input_size=416, CLASSES=TRAIN_CLASSES)
            yolo.load_weights(checkpoint)  # use custom weights

    elif YOLO_FRAMEWORK == "trt": # TensorRT detection
        saved_model_loaded = tf.saved_model.load(YOLO_CUSTOM_WEIGHTS, tags=[tag_constants.SERVING])
        signature_keys = list(saved_model_loaded.signatures.keys())
        yolo = saved_model_loaded.signatures['serving_default']

    return yolo

任何帮助将不胜感激。

标签： pythontensorflowcomputer-visionyolo

python - Yolov3.weights 未完全读取错误

问题描述

解决方案

推荐阅读