python 提取视频帧保存为图像-目标检测

摘要

针对视频数据提取想要的间隔帧，在另存为图像进行保存。

视频数据

本地视频数据如下所示。part1为mp4视频数据，part2位txt标准信息。


具体信息如下，开始是第几帧，之后是目标检测的标注信息。

代码提取图像和生成coco.json

import os
import cv2
import json
save_path="D:/racedata/2022boat\end\img/"
annotations_info = {'images': [], 'annotations': [], 'categories': []}
categories_map = {'lighthouse': 1,'sailboat':2, 'buoy': 3,'railbar':4,'cargoship': 5, 'navalvessels': 6,'passengership':7,'dock':8,'submarine':9,'fishingboat': 10, 'freighter': 11, 'pier': 12, 'tubeline': 13}
for key in categories_map:
    categoriy_info = {"id":categories_map[key], "name":key}
    annotations_info['categories'].append(categoriy_info)

file=os.listdir("D:/racedata/2022boat\end\part1")
file=[img[:-4] for img in file]
def change_txt(file_txt):
    data = open(file_txt, 'r')
    all_data = []
    for line in data.readlines():
        line = line.strip()
        line = line.split(",")
        all_data.append((line))
    return all_data


img_i=0
ann_id=1
for id in file:
    print(id)
    file_vid = "D:/racedata/2022boat\end/part1/"+id+".mp4"
    file_txt = "D:/racedata/2022boat\end/part2/"+id+"_gt.txt"
    all_data =change_txt(file_txt)
    len_data = len(all_data)
    video = cv2.VideoCapture(file_vid)
    frameCount = video.get(cv2.CAP_PROP_FRAME_COUNT)
    success, frame = video.read()
    index=0
    while success:
        if index >= len_data:
            break
        anno_line = all_data[index]
        zhen=anno_line[0]
        if int(zhen) % 50 ==0 and  int(zhen)==index:
            num = anno_line[1]
            data_line = anno_line[2:]
            #cv2.imwrite(save_path + id+"_"+zhen+".jpg", frame)
            image_file_name=id+"_"+zhen+".jpg"
            height, width, _ = frame.shape
            image_info = {'file_name': image_file_name, 'id': img_i + 1,
                          'height': height, 'width': width}
            img_i += 1
            annotations_info['images'].append(image_info)
            k=0
            for i in range(int(num)):
                single_line = data_line[k:k + 5]
                x1 = int(single_line[0])
                y1 = int(single_line[1])
                x2 = int(single_line[2])
                y2 = int(single_line[3])
                if x2<x1 and y2<y1:
                    t1=x1
                    t2=y1
                    x1=x2
                    y1=y2
                    x2=t1
                    y2=t2
                name=single_line[4]
                if x2 == width:
                    x2 -= 1
                if y2 == height:
                    y2 -= 1
                x, y = x1, y1
                w, h = x2 - x1 + 1, y2 - y1 + 1
                category_id = categories_map[name]
                area = w * h
                if (area < 5):
                    continue
                annotation_info = {"id": ann_id, "image_id":  img_i, "bbox": [x, y, w, h], "category_id": category_id,
                                   "area": area, "iscrowd": 0,"segmentation": [10,20,30,40]}
                annotations_info['annotations'].append(annotation_info)
                ann_id += 1
                k=k+5
        success, frame = video.read()
        index+=1

with  open("D:/racedata/2022boat\end"+'/instances_train_50_init_change2.json', 'w')  as f:
    json.dump(annotations_info, f, indent=4)

print('---整理后的标注文件---')
print('所有图片的数量：',  len(annotations_info['images']))
print('所有标注的数量：',  len(annotations_info['annotations']))
print('所有类别的数量：',  len(annotations_info['categories']))

每个50帧提取一张图像。

验证阶段

import cv2
import os
from tqdm import tqdm
import json
import matplotlib.pyplot as plt
ann_path = 'D:/racedata/2022boat\end/instances_train.json'
img_path = 'D:/racedata/2022boat\end/img/'
save_path = 'D:/racedata/2022boat/end/show/' # the path of saveing image with annotated bboxes
with open(ann_path,'r') as f:
    ann = json.load(f)
for ann_img in tqdm(ann['images']):
    img = cv2.imread(img_path + ann_img['file_name'])
    img_id = ann_img['id']
    for ann_ann in ann['annotations']:
        if ann_ann['image_id'] == img_id:
            x1 = int(ann_ann['bbox'][0])
            y1 = int(ann_ann['bbox'][1])
            x2 = int(ann_ann['bbox'][0] + ann_ann['bbox'][2])
            y2 = int(ann_ann['bbox'][1] + ann_ann['bbox'][3])
            img = cv2.rectangle(img, (x1,y1), (x2,y2), (255,0,0), 8)
            for cat in ann['categories']:
                if cat['id'] == ann_ann['category_id']:
                    catname = cat['name']
                    break
            txt = catname
            img = cv2.putText(img, txt, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 4)
    cv2.imwrite(save_path + ann_img['file_name'], img)

通过生成的json对图像数据进行标注画框来检查是否有误。

预测json生成标注对比

import cv2
import os
from tqdm import tqdm
import json
from itertools import groupby
file_name_change={""}
ann_path = './test_part1.json'
img_path = './img/'
save_path = './test_show_img/'
with open(ann_path,'r') as f:
    ann = json.load(f)
name_dict={1:"0001",2:"0002",3:"0003",4:"0004"}
cate_name={1:'freighter',2:'buoy',3:'lighthouse',4:'sailboat',5:'pier',6:'tubeline'}
def change_name(k):
    t=k/10000
    t=name_dict[int(t)]
    tail=k%10000
    img_name=t+"_"+str(tail)+".jpg"
    return img_name
for k, items in groupby(ann, key=lambda x:x["image_id"]):
    img_name=change_name(k)
    print(img_name)
    img = cv2.imread(img_path +img_name)
    for i in items:
        if float(i['score'])<0.3:
            continue
        x1 = int(i['bbox'][0])
        y1 = int(i['bbox'][1])
        x2 = int(i['bbox'][0] + i['bbox'][2])
        y2 = int(i['bbox'][1] + i['bbox'][3])
        img = cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 8)
        txt = cate_name[i['category_id']]
        img = cv2.putText(img, txt, (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 4)
    cv2.imwrite(save_path + img_name, img)

mmdet生成图像预测的json，在生成标注图像查看效果。