0822上传模型标注代码
parent
8d57b523d8
commit
60a14775b7
@ -0,0 +1,95 @@
|
|||||||
|
|
||||||
|
from xml.etree.ElementTree import ElementTree, Element
|
||||||
|
|
||||||
|
|
||||||
|
# xml换行
|
||||||
|
|
||||||
|
|
||||||
|
def indent(elem, level=0):
|
||||||
|
i = "\n" + level*"\t"
|
||||||
|
if len(elem):
|
||||||
|
if not elem.text or not elem.text.strip():
|
||||||
|
elem.text = i + "\t"
|
||||||
|
if not elem.tail or not elem.tail.strip():
|
||||||
|
elem.tail = i
|
||||||
|
for elem in elem:
|
||||||
|
indent(elem, level+1)
|
||||||
|
if not elem.tail or not elem.tail.strip():
|
||||||
|
elem.tail = i
|
||||||
|
else:
|
||||||
|
if level and (not elem.tail or not elem.tail.strip()):
|
||||||
|
elem.tail = i
|
||||||
|
|
||||||
|
|
||||||
|
def add_xml(inforsDict,xmlFilePath):
|
||||||
|
|
||||||
|
result = inforsDict
|
||||||
|
|
||||||
|
for re in result:
|
||||||
|
# if re['score'] > 0.5:
|
||||||
|
|
||||||
|
# 获得标注信息
|
||||||
|
ObjName = list(re.keys())[0]
|
||||||
|
xmin = int(list(re.values())[0][0])
|
||||||
|
ymin = int(list(re.values())[0][1])
|
||||||
|
xmax = int(list(re.values())[0][2])
|
||||||
|
ymax = int(list(re.values())[0][3])
|
||||||
|
# xmax = xmin + r
|
||||||
|
# ymax = ymin + z
|
||||||
|
|
||||||
|
#if ObjName == 'person':
|
||||||
|
|
||||||
|
tree = ElementTree()
|
||||||
|
tree.parse(xmlFilePath)
|
||||||
|
|
||||||
|
# 得到根目录
|
||||||
|
root = tree.getroot()
|
||||||
|
|
||||||
|
# 创建一级目录
|
||||||
|
elementOjb = Element('object')
|
||||||
|
|
||||||
|
elementBox = Element('bndbox')
|
||||||
|
|
||||||
|
# 创建二级目录
|
||||||
|
one = Element('name')
|
||||||
|
one.text = ObjName # 二级目录的值 #结果展示:<id>1</id>
|
||||||
|
elementOjb.append(one) # 将二级目录加到一级目录里
|
||||||
|
|
||||||
|
two = Element('pose')
|
||||||
|
two.text = "Unspecified"
|
||||||
|
elementOjb.append(two)
|
||||||
|
|
||||||
|
three = Element('truncated')
|
||||||
|
three.text = "0"
|
||||||
|
elementOjb.append(three)
|
||||||
|
|
||||||
|
four = Element('difficult')
|
||||||
|
four.text = "0"
|
||||||
|
elementOjb.append(four)
|
||||||
|
|
||||||
|
five = Element('xmin')
|
||||||
|
five.text = str(xmin)
|
||||||
|
elementBox.append(five)
|
||||||
|
|
||||||
|
six = Element('xmax')
|
||||||
|
six.text = str(xmax)
|
||||||
|
elementBox.append(six)
|
||||||
|
|
||||||
|
seven = Element('ymin')
|
||||||
|
seven.text = str(ymin)
|
||||||
|
elementBox.append(seven)
|
||||||
|
|
||||||
|
eight = Element('ymax')
|
||||||
|
eight.text = str(ymax)
|
||||||
|
elementBox.append(eight)
|
||||||
|
|
||||||
|
# 将一级目录加到根目录里
|
||||||
|
elementOjb.append(elementBox)
|
||||||
|
root.append(elementOjb)
|
||||||
|
# 换行缩进
|
||||||
|
indent(elementOjb)
|
||||||
|
indent(elementBox)
|
||||||
|
# 让结果保存进文件就可以了
|
||||||
|
tree.write(xmlFilePath, encoding='utf-8', xml_declaration=True)
|
||||||
|
|
||||||
|
|
@ -0,0 +1,152 @@
|
|||||||
|
# Copyright 2020 The MediaPipe Authors.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
"""MediaPipe solution drawing utils."""
|
||||||
|
|
||||||
|
import math
|
||||||
|
from typing import List, Mapping, Optional, Tuple, Union
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import dataclasses
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from mediapipe.framework.formats import detection_pb2
|
||||||
|
from mediapipe.framework.formats import location_data_pb2
|
||||||
|
from mediapipe.framework.formats import landmark_pb2
|
||||||
|
|
||||||
|
_PRESENCE_THRESHOLD = 0.5
|
||||||
|
_VISIBILITY_THRESHOLD = 0.5
|
||||||
|
_BGR_CHANNELS = 3
|
||||||
|
|
||||||
|
WHITE_COLOR = (224, 224, 224)
|
||||||
|
BLACK_COLOR = (0, 0, 0)
|
||||||
|
RED_COLOR = (0, 0, 255)
|
||||||
|
GREEN_COLOR = (0, 128, 0)
|
||||||
|
BLUE_COLOR = (255, 0, 0)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclasses.dataclass
|
||||||
|
class DrawingSpec:
|
||||||
|
# Color for drawing the annotation. Default to the white color.
|
||||||
|
color: Tuple[int, int, int] = WHITE_COLOR
|
||||||
|
# Thickness for drawing the annotation. Default to 2 pixels.
|
||||||
|
thickness: int = 2
|
||||||
|
# Circle radius. Default to 2 pixels.
|
||||||
|
circle_radius: int = 2
|
||||||
|
|
||||||
|
|
||||||
|
def _normalized_to_pixel_coordinates(
|
||||||
|
normalized_x: float, normalized_y: float, image_width: int,
|
||||||
|
image_height: int) -> Union[None, Tuple[int, int]]:
|
||||||
|
"""Converts normalized value pair to pixel coordinates."""
|
||||||
|
|
||||||
|
# Checks if the float value is between 0 and 1.
|
||||||
|
def is_valid_normalized_value(value: float) -> bool:
|
||||||
|
return (value > 0 or math.isclose(0, value)) and (value < 1 or
|
||||||
|
math.isclose(1, value))
|
||||||
|
|
||||||
|
if not (is_valid_normalized_value(normalized_x) and
|
||||||
|
is_valid_normalized_value(normalized_y)):
|
||||||
|
# TODO: Draw coordinates even if it's outside of the image bounds.
|
||||||
|
return None
|
||||||
|
x_px = min(math.floor(normalized_x * image_width), image_width - 1)
|
||||||
|
y_px = min(math.floor(normalized_y * image_height), image_height - 1)
|
||||||
|
return x_px, y_px
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def draw_landmarks(
|
||||||
|
image: np.ndarray,
|
||||||
|
landmark_list: landmark_pb2.NormalizedLandmarkList,
|
||||||
|
connections: Optional[List[Tuple[int, int]]] = None):
|
||||||
|
"""Draws the landmarks and the connections on the image.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: A three channel BGR image represented as numpy ndarray.
|
||||||
|
landmark_list: A normalized landmark list proto message to be annotated on
|
||||||
|
the image.
|
||||||
|
connections: A list of landmark index tuples that specifies how landmarks to
|
||||||
|
be connected in the drawing.
|
||||||
|
landmark_drawing_spec: Either a DrawingSpec object or a mapping from hand
|
||||||
|
landmarks to the DrawingSpecs that specifies the landmarks' drawing
|
||||||
|
settings such as color, line thickness, and circle radius. If this
|
||||||
|
argument is explicitly set to None, no landmarks will be drawn.
|
||||||
|
connection_drawing_spec: Either a DrawingSpec object or a mapping from hand
|
||||||
|
connections to the DrawingSpecs that specifies the connections' drawing
|
||||||
|
settings such as color and line thickness. If this argument is explicitly
|
||||||
|
set to None, no landmark connections will be drawn.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If one of the followings:
|
||||||
|
a) If the input image is not three channel BGR.
|
||||||
|
b) If any connetions contain invalid landmark index.
|
||||||
|
"""
|
||||||
|
if not landmark_list:
|
||||||
|
return
|
||||||
|
if image.shape[2] != _BGR_CHANNELS:
|
||||||
|
raise ValueError('Input image must contain three channel bgr data.')
|
||||||
|
image_rows, image_cols, _ = image.shape
|
||||||
|
|
||||||
|
# 所有的点转换成坐标的字典
|
||||||
|
idx_to_coordinates = {}
|
||||||
|
for idx, landmark in enumerate(landmark_list.landmark):
|
||||||
|
# print('landmark:',landmark)
|
||||||
|
if ((landmark.HasField('visibility') and
|
||||||
|
landmark.visibility < _VISIBILITY_THRESHOLD) or
|
||||||
|
(landmark.HasField('presence') and
|
||||||
|
landmark.presence < _PRESENCE_THRESHOLD)):
|
||||||
|
continue
|
||||||
|
landmark_px = _normalized_to_pixel_coordinates(landmark.x, landmark.y,
|
||||||
|
image_cols, image_rows)
|
||||||
|
# print('landmark_px:',landmark_px)
|
||||||
|
if landmark_px:
|
||||||
|
idx_to_coordinates[idx] = landmark_px
|
||||||
|
|
||||||
|
|
||||||
|
if connections:
|
||||||
|
num_landmarks = len(landmark_list.landmark)
|
||||||
|
# print('connections:',connections)
|
||||||
|
|
||||||
|
# Draws the connections if the start and end landmarks are both visible.
|
||||||
|
|
||||||
|
start_list = []
|
||||||
|
end_list = []
|
||||||
|
for connection in connections:
|
||||||
|
# print(connection)
|
||||||
|
|
||||||
|
start_idx = connection[0]
|
||||||
|
end_idx = connection[1]
|
||||||
|
|
||||||
|
start_list.append(start_idx)
|
||||||
|
end_list.append(end_idx)
|
||||||
|
|
||||||
|
|
||||||
|
point_list = []
|
||||||
|
for point_idx in end_list:
|
||||||
|
|
||||||
|
# if point_idx not in start_list:
|
||||||
|
|
||||||
|
# print(point_idx)
|
||||||
|
point_list.append(point_idx)
|
||||||
|
|
||||||
|
|
||||||
|
point_axis_list = []
|
||||||
|
for point in point_list:
|
||||||
|
|
||||||
|
if point in list(idx_to_coordinates.keys()):
|
||||||
|
point_axis_list.append(idx_to_coordinates[point])
|
||||||
|
|
||||||
|
|
||||||
|
return point_axis_list
|
||||||
|
|
@ -0,0 +1,181 @@
|
|||||||
|
import cv2
|
||||||
|
import os
|
||||||
|
import mediapipe as mp
|
||||||
|
from ultralytics import YOLO
|
||||||
|
from personDet import analysis_yolov8
|
||||||
|
import tools_function
|
||||||
|
from holisticDet import MediapipeProcess
|
||||||
|
from add_xml import add_xml
|
||||||
|
from create_xml import create_xml
|
||||||
|
import queue
|
||||||
|
|
||||||
|
|
||||||
|
class DetProcess():
|
||||||
|
|
||||||
|
def __init__(self,person_det_model,hand_det_model):
|
||||||
|
|
||||||
|
self.person_det_model = person_det_model
|
||||||
|
self.hand_det_model = hand_det_model
|
||||||
|
|
||||||
|
|
||||||
|
def get_person_cut(self,frame,det_dict,imgsize):
|
||||||
|
|
||||||
|
# person_det_dict = [perdict for perdict in det_dict if list(perdict.keys())[0] == 'person']
|
||||||
|
|
||||||
|
# print('person_det_dict:',det_dict)
|
||||||
|
person_list = tools_function.get_dict_values(det_dict)
|
||||||
|
|
||||||
|
# 坐标参数修正
|
||||||
|
person_bbox_list = tools_function.para_list_correction(images_size=imgsize,bbox_list=person_list,dertpara=5)
|
||||||
|
|
||||||
|
frame_list = []
|
||||||
|
for per_bbox in person_bbox_list:
|
||||||
|
|
||||||
|
# 裁剪后人的图片
|
||||||
|
person_cut_frame = tools_function.img_cut(frame=frame,bbox=per_bbox)
|
||||||
|
|
||||||
|
frame_cut_dict = {tuple(per_bbox):person_cut_frame}
|
||||||
|
|
||||||
|
frame_list.append(frame_cut_dict)
|
||||||
|
|
||||||
|
return frame_list
|
||||||
|
|
||||||
|
def hand_det(self,person_cut_frame,per_bbox):
|
||||||
|
|
||||||
|
# print('11111')
|
||||||
|
hands_result = MediapipeProcess.mediapipe_det(image=person_cut_frame,holistic=self.hand_det_model)
|
||||||
|
hands_result_dict = MediapipeProcess.get_analysis_result(image=person_cut_frame,results=hands_result)
|
||||||
|
hands_list = tools_function.select_list(hands_result_dict['hand_bbox'])
|
||||||
|
|
||||||
|
# print('hands_list:',hands_list)
|
||||||
|
imgsize2 = person_cut_frame.shape
|
||||||
|
|
||||||
|
# 手部坐标修正
|
||||||
|
hands_bbox_list = tools_function.para_list_correction(images_size=imgsize2,bbox_list=hands_list,dertpara=5)
|
||||||
|
|
||||||
|
# print('hands_bbox_list:',hands_bbox_list)
|
||||||
|
|
||||||
|
hand_bbox_list = []
|
||||||
|
for hand in hands_bbox_list:
|
||||||
|
|
||||||
|
hands_result_list = tools_function.change_bbox(bbox_person=[per_bbox[0],per_bbox[1]],bbox_hand=hand)
|
||||||
|
|
||||||
|
# print('hands_result_list:',hands_result_list)
|
||||||
|
|
||||||
|
re_dict = {'hands':hands_result_list}
|
||||||
|
|
||||||
|
hand_bbox_list.append(re_dict)
|
||||||
|
|
||||||
|
# hands_result_original_dict = {'results':hand_bbox_list}
|
||||||
|
|
||||||
|
# print(hands_result_original_dict)
|
||||||
|
|
||||||
|
return hand_bbox_list
|
||||||
|
|
||||||
|
def save_annotations_xml(self,xml_save_file,save_infors,images):
|
||||||
|
|
||||||
|
# images = save_infors['images']
|
||||||
|
results = save_infors
|
||||||
|
img = os.path.basename(images)
|
||||||
|
img_frame = cv2.imread(images)
|
||||||
|
xml_save_path = os.path.join(xml_save_file,img.split('.')[0] + '.xml')
|
||||||
|
w,h,d = img_frame.shape
|
||||||
|
img_shape = (w,h,d,img)
|
||||||
|
|
||||||
|
if os.path.isfile(xml_save_path):
|
||||||
|
|
||||||
|
add_labels = add_xml(inforsDict=results,xmlFilePath=xml_save_path)
|
||||||
|
else:
|
||||||
|
create_new = create_xml(boxs=results,img_shape=img_shape,xml_path=xml_save_path)
|
||||||
|
|
||||||
|
|
||||||
|
def person_cut_process(self,images,img_save_files):
|
||||||
|
|
||||||
|
|
||||||
|
frame = cv2.imread(images)
|
||||||
|
imgsize = frame.shape
|
||||||
|
labels_name_list = ['person']
|
||||||
|
|
||||||
|
per_det_dict = analysis_yolov8(frame=frame,
|
||||||
|
model_coco=self.person_det_model,
|
||||||
|
labels_names=labels_name_list,
|
||||||
|
confidence_set=0.2)
|
||||||
|
|
||||||
|
per_frame_cut = self.get_person_cut(frame=frame,det_dict=per_det_dict,imgsize=imgsize)
|
||||||
|
per_frame_list = [value for dictionary in per_frame_cut for value in dictionary.values()]
|
||||||
|
|
||||||
|
# print('per_frame_list:',per_frame_list)
|
||||||
|
|
||||||
|
for id_num,cut_frame in enumerate(per_frame_list):
|
||||||
|
|
||||||
|
cut_frame_save = tools_function.img_write(frame=cut_frame,img_file=images,id_num=id_num,save_file=img_save_files)
|
||||||
|
|
||||||
|
def hands_det_process(self,images,xml_save_file):
|
||||||
|
|
||||||
|
frame = cv2.imread(images)
|
||||||
|
imgsize = frame.shape
|
||||||
|
labels_name_list = ['person']
|
||||||
|
|
||||||
|
per_det_dict = analysis_yolov8(frame=frame,
|
||||||
|
model_coco=self.person_det_model,
|
||||||
|
labels_names=labels_name_list,
|
||||||
|
confidence_set=0.2)
|
||||||
|
|
||||||
|
per_frame_cut = self.get_person_cut(frame=frame,det_dict=per_det_dict,imgsize=imgsize)
|
||||||
|
|
||||||
|
for frame_dict in per_frame_cut:
|
||||||
|
|
||||||
|
per_bbox = list(frame_dict.keys())[0]
|
||||||
|
person_cut_frame = list(frame_dict.values())[0]
|
||||||
|
|
||||||
|
hands_dict = self.hand_det(person_cut_frame=person_cut_frame,per_bbox=per_bbox)
|
||||||
|
|
||||||
|
self.save_annotations_xml(xml_save_file=xml_save_file,save_infors=hands_dict,images=images)
|
||||||
|
|
||||||
|
|
||||||
|
# person_det = self.detect_yolo(images_path=images,labels_name_list=labels_name_list)
|
||||||
|
|
||||||
|
def det_process(self,images,xml_save_file):
|
||||||
|
|
||||||
|
frame = cv2.imread(images)
|
||||||
|
imgsize = frame.shape
|
||||||
|
labels_name_list = ['cell phone','mouse','keyboard']
|
||||||
|
|
||||||
|
per_det_dict = analysis_yolov8(frame=frame,
|
||||||
|
model_coco=self.person_det_model,
|
||||||
|
labels_names=labels_name_list,
|
||||||
|
confidence_set=0.2)
|
||||||
|
|
||||||
|
self.save_annotations_xml(xml_save_file=xml_save_file,save_infors=per_det_dict,images=images)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
|
||||||
|
images_files = 'images'
|
||||||
|
images_list = tools_function.get_path_list(images_files)
|
||||||
|
img_save_files = 'images_cut'
|
||||||
|
xml_save_file = 'annotations'
|
||||||
|
|
||||||
|
|
||||||
|
# 初始化目标检测
|
||||||
|
person_model = YOLO("model_files/yolov8x.pt")
|
||||||
|
|
||||||
|
# 初始化mediapipe
|
||||||
|
mp_holistic = mp.solutions.holistic
|
||||||
|
holistic = mp_holistic.Holistic(min_detection_confidence=0.1,min_tracking_confidence=0.1)
|
||||||
|
|
||||||
|
Det = DetProcess(person_det_model=person_model,hand_det_model=holistic)
|
||||||
|
|
||||||
|
for images in images_list:
|
||||||
|
|
||||||
|
# Det.person_cut_process(images=images,img_save_files=img_save_files)
|
||||||
|
Det.hands_det_process(images=images,xml_save_file=xml_save_file)
|
||||||
|
Det.det_process(images=images,xml_save_file=xml_save_file)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue