Python YOLO_V1 Examples

Example #1

#---------------step1:Dataset-------------------
import torch
from YOLO_V1_DataSet import VOCDataSet
dataSet = VOCDataSet(imgs_dir="./VOC2007/Train/JPEGImages",
                     annotations_dir="./VOC2007/Train/Annotations",
                     ClassesFile="./VOC2007/Train/class.data")
from torch.utils.data import DataLoader
#dataLoader = DataLoader(dataSet,batch_size=64,shuffle=True,num_workers=0)

#---------------step2:Model-------------------
from YOLO_V1_Model import YOLO_V1
YOLO = YOLO_V1().cuda(device=0)
YOLO.initialize_weights()

#---------------step3:LossFunction-------------------
from YOLO_V1_LossFunction import YOLO_V1_Loss
loss_function = YOLO_V1_Loss().cuda(device=0)

#---------------step4:Optimizer-------------------
import torch.optim as optim
optimizer_Adam = optim.Adam(YOLO.parameters(), lr=1e-4, weight_decay=0.0005)
#使用余弦退火动态调整学习率
#lr_reduce_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer_Adam , T_max=20, eta_min=1e-4, last_epoch=-1)
#lr_reduce_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer=optimizer_Adam, T_0=2, T_mult=2)
#--------------step5:Tensorboard Feature Map------------
from tensorboardX import SummaryWriter
import torchvision.utils as vutils
import torch.nn as nn


def feature_map_visualize(img_data, writer):

Example #2

File: YOLO_V1_Test.py Project: ProgrammerZhujinming/YOLO_V1_GPU

# 网络加载
from YOLO_V1_Model import YOLO_V1

YOLO = YOLO_V1()
import torch

YOLO.load_state_dict(torch.load('./YOLO_V1_300.pth'))
# 类别与索引转换
IndexToClassName = {}
with open("./VOC2007/Train/class.data", "r") as f:
    index = 0
    for line in f:
        IndexToClassName[index] = line
        index = index + 1


def iou(box_one, box_two):
    LX = max(box_one[0], box_two[0])
    LY = max(box_one[1], box_two[1])
    RX = min(box_one[2], box_two[2])
    RY = min(box_one[3], box_two[3])
    if LX >= RX or LY >= RY:
        return 0
    return (RX - LX) * (RY - LY) / ((box_one[2] - box_one[0]) *
                                    (box_one[3] - box_one[1]) +
                                    (box_two[2] - box_two[0]) *
                                    (box_two[3] - box_two[1]))


import numpy as np

Example #3

import torch
from YOLO_V1_Model import YOLO_V1

weight_file_name = "YOLO_V1_2400.pth"

model = YOLO_V1().cuda(device=1)
model.load_state_dict(torch.load(weight_file_name))
torch.save(model.state_dict(),
           weight_file_name,
           _use_new_zipfile_serialization=False)

Example #4

#------step1: 导入网络------
from YOLO_V1_Model import YOLO_V1
YOLO = YOLO_V1().cuda()

#------step:2 读取权重文件------
import torch
weight_file_name = "./Model_Train/YOLO_V1_500.pth"
YOLO.load_state_dict(torch.load(weight_file_name)["model"])
YOLO.eval()

#------step:3 类别索引与类别名的映射------
class_file_name = "./VOC2007/Train/class.data"
class_index_Name = {}
classIndex = 0
with open(class_file_name, 'r') as f:
    for line in f:
        line = line.replace('\n', '')
        class_index_Name[classIndex] = line  # 根据类别名制作索引
        classIndex = classIndex + 1

#------step:4 NMS算法处理输出结果------


def iou(box_one, box_two):
    LX = max(box_one[0], box_two[0])
    LY = max(box_one[1], box_two[1])
    RX = min(box_one[2], box_two[2])
    RY = min(box_one[3], box_two[3])
    if LX >= RX or LY >= RY:
        return 0
    return (RX - LX) * (RY - LY) / ((box_one[2] - box_one[0]) *