python目標檢測實現黑花屏分類任務示例

背景

視頻幀的黑、花屏的檢測是視頻質量檢測中比較重要的一部分,傳統做法是由測試人員通過肉眼來判斷視頻中是否有黑、花屏的現象,這種方式不僅耗費人力且效率較低。

為瞭進一步節省人力、提高效率,一種自動的檢測方法是大傢所期待的。目前,通過分類網絡模型對視頻幀進行分類來自動檢測是否有黑、花屏是比較可行且高效的。

然而,在項目過程中,視頻幀數據的收集比較困難,數據量較少,部分花屏和正常屏之間差異不夠明顯,導致常用的分類算法難以滿足項目對分類準確度的要求。

因此本文嘗試瞭一種利用目標檢測算法實現分類的方式,幫助改善單純的分類的算法效果不夠理想的問題。

核心技術與架構圖

一般分類任務的流程如下圖,首先需要收集數據,構成數據集;

並為每一類數據定義一個類型標簽,例如:0、1、2;再選擇一個合適的分類網絡進行分類模型的訓練,圖像分類的網絡有很多,常見的有 VggNet, ResNet,DenseNet 等;

最後用訓練好的模型對新的數據進行預測,輸出新數據的類別。

目標檢測任務的流程不同於分類任務,其在定義類別標簽的時候還需要對目標位置進行標註;

目標檢測的方法也有很多,例如 Fast R-CNN, SSD,YOLO 等;

模型訓練的中間過程也比分類模型要復雜,其輸出一般為目標的位置、目標置信度以及分類結果。

由於分類算法依賴於一定量的數據,在項目實踐中,數據量較少或圖像類間差異較小時,傳統分類算法效果不一定能滿足項目需求。這時,不妨考慮用目標檢測的方式來做 ‘分類’。

接下來以 Yolov5 為例來介紹如何將目標檢測框架用於實現單純的分類任務。

技術實現

除瞭分類之外,目標檢測還可以從自然圖像中的大量預定義類別中識別出目標實例的位置。

大傢可能會考慮目標檢測模型用於分類是不是過於繁瑣或者用目標檢測框架來做單純的分類對代碼的修改比較復雜。

這裡,我們將用一種非常簡單的方式直接在數據標註和輸出內容上稍作修改就能實現單純的分類瞭。接下來將介紹一下具體實現方法:

1.數據的標註

實現目標檢測時,需要對數據中的目標進行標註,這一過程是十分繁瑣的。但在用於純粹的分類上可以將這一繁瑣過程簡單化,無需手動標註,直接將整張圖作為我們的目標,目標中心也就是圖像的中心點。

隻需讀取整張圖像,獲得其長、寬以及中心點的坐標就可以完成標註瞭。並定義好類別標簽,正常屏為 0,花屏為:1,黑屏為 2。具體實現如下:

OBJECT_DICT = {"Normalscreen": 0, "Colorfulscreen": 1, "Blackscreen": 2}
def parse_json_file(image_path):
    imageName = os.path.basename(image_path).split('.')[0]
    img = cv2.imread(image_path)
    size = img.shape
    label = image_path.split('/')[4].split('\\')[0]
    label = OBJECT_DICT.get(label)
    imageWidth = size[0]
    imageHeight = size[1]
    label_dict = {}
    xmin, ymin = (0, 0)
    xmax, ymax = (imageWidth, imageHeight)
    xcenter = (xmin + xmax) / 2
    xcenter = xcenter / float(imageWidth)
    ycenter = (ymin + ymax) / 2
    ycenter = ycenter / float(imageHeight)
    width = ((xmax - xmin) / float(imageWidth))
    heigt = ((ymax - ymin) / float(imageHeight))
    label_dict.update({label: [str(xcenter), str(ycenter), str(width), str(heigt)]})
    label_dict = sorted(label_dict.items(), key=lambda x: x[0])
    return imageName, label_dict

2.訓練過程

該過程與目標檢測的訓練過程一致,不需要進行大的修改,隻需要根據數據集的特性對參數進行調整。

# 加載數據,獲取訓練集、測試集圖片路徑
with open(opt.data) as f:
    data_dict = yaml.load(f, Loader=yaml.FullLoader)  
    with torch_distributed_zero_first(rank):
        check_dataset(data_dict) 
train_path = data_dict['train']
test_path = data_dict['val']
Number_class, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) 
# 創建模型
model = Model(opt.cfg, ch=3, nc=Number_class).to(device)
# 學習率的設置
lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf']  
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
# 訓練
for epoch in range(start_epoch, epochs):  
model.train()

3.損失的計算

損失由三部分組成,邊框損失,目標損失,分類損失,具體如下:

def compute_loss(p, targets, model):
    device = targets.device
    loss_cls, loss_box, loss_obj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
    tcls, tbox, indices, anchors = build_targets(p, targets, model) 
h = model.hyp
    # 定義損失函數
    BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['cls_pw']])).to(device)
    BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['obj_pw']])).to(device)
    cp, cn = smooth_BCE(eps=0.0)
    # 損失
    nt = 0  
    np = len(p) 
    balance = [4.0, 1.0, 0.4] if np == 3 else [4.0, 1.0, 0.4, 0.1] 
for i, pi in enumerate(p): 
        image, anchor, gridy, gridx = indices[i]  
        tobj = torch.zeros_like(pi[..., 0], device=device) 
        n = image.shape[0]  
        if n:
            nt += n  # 計算目標
            ps = pi[anchor, image, gridy, gridx]
            pxy = ps[:, :2].sigmoid() * 2. - 0.5
            pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
            predicted_box = torch.cat((pxy, pwh), 1).to(device)                    giou = bbox_iou(predicted_box.T, tbox[i], x1y1x2y2=False, CIoU=True)                 
       loss_box += (1.0 - giou).mean() 
            tobj[image, anchor, gridy, gridx] = (1.0 - model.gr) + model.gr *   giou.detach().clamp(0).type(tobj.dtype) 
            if model.nc > 1:
                t = torch.full_like(ps[:, 5:], cn, device=device)
                t[range(n), tcls[i]] = cp
                loss_cls += BCEcls(ps[:, 5:], t)  
        loss_obj += BCEobj(pi[..., 4], tobj) * balance[i]  
    s = 3 / np
    loss_box *= h['giou'] * s
    loss_obj *= h['obj'] * s * (1.4 if np == 4 else 1.)
    loss_cls *= h['cls'] * s
    bs = tobj.shape[0]
    loss = loss_box + loss_obj + loss_cls
    return loss * bs, torch.cat((loss_box, loss_obj, loss_cls, loss)).detach()

4.對輸出內容的處理

進行預測時,會得到所有檢測到的目標的位置(x,y,w,h),objectness 置信度和分類結果。由於最終目的是對整張圖進行分類,可以忽略位置信息,重點考慮置信度和分類結果:將檢測到的目標類別作為分類結果,如果同時檢測出多個目標,可以將置信度最大的目標的類別作為分類結果。代碼如下:

def detect(opt,img):
    out, source, weights, view_img, save_txt, imgsz = \
        opt.output, img, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    device = select_device(opt.device)
    half = device.type != 'cpu'
    model = experimental.attempt_load(weights, map_location=device)
    imgsz = check_img_size(imgsz, s=model.stride.max())
    if half:
        model.half()
    img = letterbox(img)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1)
    img = np.ascontiguousarray(img)
    img_warm = torch.zeros((1, 3, imgsz, imgsz), device=device)
    _ = model(img_warm.half() if half else img_warm) if device.type != 'cpu' else None 
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()
    img /= 255.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    pred = model(img, augment=opt.augment)[0]
    # 應用非極大值抑制
    pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
    # 處理檢測的結果
    for i, det in enumerate(pred): 
        if det is not None and len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape).round()
            all_conf = det[:, 4]
            if len(det[:, -1]) > 1:
                ind = torch.max(all_conf, 0)[1]
                c = torch.take(det[:, -1], ind)
detect_class = int(c)
            else:
                for c in det[:, -1]:
                    detect_class = int(c)
            return detect_class

效果展示

為瞭將視頻幀進行黑、花屏分類,測試人員根據經驗將屏幕分為正常屏(200 張)、花屏(200 張)和黑屏(200 張)三類,其中正常屏幕標簽為 0,花屏的標簽為 1,黑屏的標簽為 2。

為瞭進一步說明該方法的有效性,我們將基於 Yolov5 的 ‘分類’ 效果與 ResNet 分類效果做瞭對比。根據測試人員對 ResNet 分類效果的反饋來看,ResNet 模型容易將正常屏與花屏錯誤分類,例如,下圖被測試人員定義為正常屏:

ResNet 的分類結果為 1,即為花屏,顯然,這不是我們想要的結果。

基於 Yolov5 的分類結果為 0,即為正常屏,這是我們所期待的結果。

同時,通過對一批測試數據的分類效果來看,Yolov5 的分類效果比 ResNet 的分類準確度更高,ResNet 的分類準確率為 88%,而基於 Yolov5 的分類準確率高達 97%。

總結

對於較小數據集的黑、花屏的分類問題,采用 Yolov5 來實現分類相較於 ResNet 的分類效果會更好一些。當我們在做圖像分類任務時,純粹的分類算法不能達到想要的效果時,不妨嘗試一下用目標檢測框架來分類吧!雖然過程稍微復雜一些,但可能會有不錯的效果。

目前目標檢測框架有很多,用它們完成分類任務的處理方式大致和本文所描述的類似,可以根據數據集的特征選擇合適目標檢測架構來實現分類。

本文主要介紹瞭如何將現有的目標檢測框架直接用於單純的圖像分類任務,當然,為瞭使得結構更簡潔,也可以將目標檢測中的分類網絡提取出來用於分類,更多關於python目標檢測黑花屏分類的資料請關註WalkonNet其它相關文章!

推薦閱讀: