yatman
/
Image_Gen_Server
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

gened yolo dataset

This commit is contained in:
zaqxs123456 2024-10-21 12:38:29 +08:00
parent de2b354196
commit 8fb4f8753d
5 changed files with 101 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -172,3 +172,4 @@ avg_keypoints/*
video_frames/*
video/*
new_yolo_keypoints/*
new_yolo_keypoints_bak/*

27
openpose_gen.py
View File

@ -104,7 +104,7 @@ def upload_image_circular_queue(image_path, size, unqiue_id, server_address):
return image_name
def main():
def visualize_for_fixed_dataset(json_file: str):
directory = './fixed'
json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
if not json_files:
@ -134,5 +134,30 @@ def main():
cv2.imwrite(image_path + '_' + str(i) + '.png', canvas)
def visualize_for_new_yolo_dataset():
directory = './new_yolo_keypoints'
json_files = [f for f in os.listdir(directory) if f.endswith('.json')]
if not json_files:
print("No JSON files found in the directory.")
return
json_file = os.path.join(directory, random.choice(json_files))
# create ./output directory if it does not exist
os.makedirs('output', exist_ok=True)
frames = pjf.read_new_yolo_keypoints(json_file)
for i, skeletons in enumerate(frames):
canvas = np.zeros((360, 640, 3), dtype=np.uint8)
for skeleton in skeletons:
keypoints = skeleton.keypoints
canvas = skel.draw_bodypose(canvas, keypoints, skel.yolo_coco_limbSeq, skel.yolo_coco_colors)
print(os.path.join('./output', os.path.basename(json_file) + '_' + str(i) + '.png'))
cv2.imwrite(os.path.join('./output', 'a_test' + '_' + str(i) + '.png'), canvas)
def main():
visualize_for_new_yolo_dataset()
if __name__ == '__main__':
main()

32
process_json_file.py
View File

@ -31,6 +31,34 @@ def array_json_to_Skeleton_Seqences(json_file: str) -> List[skel.Skeleton_Seqenc
skeleton_sequences[i].add_frame(skeleton)
return skeleton_sequences
def read_new_yolo_keypoints(json_file: str) -> List[List[skel.Skeleton]]:
with open(json_file, 'r') as file:
data = json.load(file)
frames = []
for frame in data:
skeletons = []
for skeleton in frame:
keypoints = [skel.Keypoint(point['x'], point['y'], point['confidence']) for point in skeleton['keypoints']]
skeletons.append(skel.Skeleton(keypoints))
frames.append(skeletons)
return frames
def add_neck_to_og_coco(folder: str):
for file in os.listdir(folder):
with open(os.path.join(folder, file), 'r') as f:
data = json.load(f)
print(f"Processing {file}")
for frame in data:
for skeleton in frame:
_skeleton = skeleton['keypoints']
# Add neck keypoint to the original COCO keypoints
if len(_skeleton) > 6:
neck = {'x': (_skeleton[6]['x'] + _skeleton[7]['x']) / 2, 'y': (_skeleton[6]['y'] + _skeleton[7]['y']) / 2, 'confidence': (_skeleton[6]['confidence'] + _skeleton[7]['confidence']) / 2}
_skeleton.insert(6, neck)
with open(os.path.join(folder, file), 'w') as f:
json.dump(data, f, indent=4)
def Skeleton_Seqences_save_to_array_json(skeleton_sequences: List[skel.Skeleton_Seqence], json_file: str):
# Ensure the directory exists
os.makedirs(os.path.dirname(json_file), exist_ok=True)
@ -287,11 +315,11 @@ def main():
video_path = './video'
video_frame_path = './video_frames'
add_neck_to_og_coco('./new_yolo_keypoints')
print("Done processing all rows.")
# print("Done processing all rows.")
# canvas = np.zeros((360, 640, 3), dtype=np.uint8)
# canvas = skel.draw_bodypose(canvas, frames[0][0].keypoints, skel.body_25_limbSeq, skel.body_25_colors)

49
skeleton_lib.py
View File

@ -19,6 +19,19 @@ coco_colors = [
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
]
yolo_coco_colors = [
[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]
]
yolo_coco_limbSeq = [
[1, 2], [1, 5], [2, 3], [3, 4],
[5, 6], [6, 7], [1, 8], [8, 9],
[9, 10], [1, 11], [11, 12], [12, 13],
[1, 0], [0, 14], [14, 16], [0, 15],
]
body_25_limbSeq = [
[1, 8], [1, 2], [1, 5], [2, 3],
[3, 4], [5, 6], [6, 7], [8, 9],
@ -176,28 +189,28 @@ def draw_bodypose(canvas: ndarray, keypoints: List[Keypoint], limbSeq, colors, x
if keypoints is None or len(keypoints) == 0:
return canvas
for (k1_index, k2_index), color in zip(limbSeq, colors):
keypoint1 = keypoints[k1_index]
keypoint2 = keypoints[k2_index]
# for (k1_index, k2_index), color in zip(limbSeq, colors):
# keypoint1 = keypoints[k1_index]
# keypoint2 = keypoints[k2_index]
if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0:
# if keypoint1 is None or keypoint1.confidence == 0:
# print(f"keypoint failed: {k1_index}")
# if keypoint2 is None or keypoint2.confidence == 0:
# print(f"keypoint failed: {k2_index}")
continue
# if keypoint1 is None or keypoint2 is None or keypoint1.confidence == 0 or keypoint2.confidence == 0 or keypoint1.x <= 0 or keypoint1.y <= 0 or keypoint2.x <= 0 or keypoint2.y <= 0:
# # if keypoint1 is None or keypoint1.confidence == 0:
# # print(f"keypoint failed: {k1_index}")
# # if keypoint2 is None or keypoint2.confidence == 0:
# # print(f"keypoint failed: {k2_index}")
# continue
Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
X = np.array([keypoint1.y, keypoint2.y]) * float(H)
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
# Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
# X = np.array([keypoint1.y, keypoint2.y]) * float(H)
# mX = np.mean(X)
# mY = np.mean(Y)
# length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
# angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
# polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth*stick_scale), int(angle), 0, 360, 1)
# cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
for keypoint, color in zip(keypoints, colors):
if keypoint is None or keypoint.confidence == 0:
if keypoint is None or keypoint.confidence == 0 or keypoint.x <= 0 or keypoint.y <= 0:
continue
x, y = keypoint.x, keypoint.y

26
yolo_main.py
View File

@ -38,7 +38,7 @@ def get_yoloed_frames(results, lerped_keypoints_path):
return frames
def process_clip(row, model):
clip_name = row['ClipName']
clip_name = row
input_video_path = f"video_frames/{clip_name}"
lerped_keypoints_path = f"./lerped_keypoints/{clip_name}.json"
output_keypoints_path = f"./new_yolo_keypoints/{clip_name}.json"
@ -47,9 +47,9 @@ def process_clip(row, model):
os.makedirs(os.path.dirname(lerped_keypoints_path), exist_ok=True)
os.makedirs(os.path.dirname(output_keypoints_path), exist_ok=True)
# # return if the file already exists
# if os.path.exists(output_keypoints_path):
# return
# return if the file already exists
if os.path.exists(output_keypoints_path):
return
results = model(input_video_path)
frames = get_yoloed_frames(results, lerped_keypoints_path)
@ -59,7 +59,7 @@ def process_clip(row, model):
json.dump(frames, f, cls=skel.Encoder, indent=4)
def process_rows_on_gpu(rows, model, device):
for _, row in rows.iterrows():
for row in rows:
for _ in range(5):
try:
process_clip(row, model)
@ -73,7 +73,11 @@ def process_rows_on_gpu(rows, model, device):
def gen_yolo_skeletons(descriptor):
num_gpus = torch.cuda.device_count()
rows_per_gpu = len(descriptor) // num_gpus
ClipName_list = descriptor['ClipName'].tolist()
finished_list = os.listdir('./new_yolo_keypoints')
ClipName_list = [clip for clip in ClipName_list if clip + '.json' not in finished_list]
rows_per_gpu = len(ClipName_list) // num_gpus
models = [YOLO("yolo11x-pose.pt").to(torch.device(f'cuda:{i}')) for i in range(num_gpus)]
@ -81,8 +85,8 @@ def gen_yolo_skeletons(descriptor):
futures = []
for i in range(num_gpus):
start_idx = i * rows_per_gpu
end_idx = (i + 1) * rows_per_gpu if i != num_gpus - 1 else len(descriptor)
gpu_rows = descriptor.iloc[start_idx:end_idx]
end_idx = (i + 1) * rows_per_gpu if i != num_gpus - 1 else len(ClipName_list)
gpu_rows = ClipName_list[start_idx:end_idx]
futures.append(executor.submit(process_rows_on_gpu, gpu_rows, models[i], torch.device(f'cuda:{i}')))
for future in concurrent.futures.as_completed(futures):
@ -91,12 +95,6 @@ def gen_yolo_skeletons(descriptor):
except Exception as e:
print(f"Error processing rows on GPU: {e}")
def gen_yolo_skeletons_single(descriptor):
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = YOLO("yolo11x-pose.pt").to(device)
process_rows_on_gpu(descriptor, model, device)
def main():
model = YOLO("yolo11x-pose.pt") # pretrained YOLO11n model
descriptor = pd.read_csv('./ClipDescriptorKaggle_processed.csv')