Newer
Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import os
import sys
import numpy as np
import time
import yaml
import cv2
import matplotlib.pyplot as plt
"""
Author: Jenö Faist, Paul Judis
"""
"""
With this File you can manual can mark a video of a 3D print to make a fail detection dataset for your printer
you can use this dataset to mix it with other Dataset like we did with our mixed Dataset for domain shift.
"""
"""
Depending on the Video you must crop it so there a no black lines in it.
You can find out the cropping by just using try and error.
(TODO) IMPLEMENT ADAPTIVE CROPING
"""
def crop_frame(image):
cromped_image = image[50:1080-200,550:1920-600,:]
return cromped_image
if __name__ == '__main__':
"""
Saving every frame would resulte in huge dataset with pretty reduandend frames so use this variable to
specifice how many frames should be skiped before saving the image in the dataset
"""
frame_save_skip = 30
"""
Specifice here the Directory of the new Dataset and what Video you using.
"""
absolutepath = os.path.dirname(__file__)
dataset_PATH = absolutepath+'/DATASETS/own_SET'
video_PATH = absolutepath+'/TRAIN_VIDEOS/rot_fail.mp4'
print("STARTING DATASET CREATER FOR 3D PRINT ERROR DETECTION")
print("--------------------")
print("Video PATH: "+video_PATH)
print("DATASET PATH: "+dataset_PATH)
print("--------------------")
print("Loading Model...")
"""
Making a Name for the new Files
"""
temp = video_PATH.replace(".", "_")
Images_Name = temp.split('/')[-1].replace("mp4","")
print("Loading VIDEO ...")
"""
Searching for the Video if the Video cant be openend or found close the programm
"""
cap = None
try:
cap = cv2.VideoCapture(video_PATH)
print("[✓] Video Loaded [✓]")
except:
print("[!?] Video couldn't be found [?!]")
exit(1)
loop_var = True
count = 0
while(loop_var):
ret, frame = cap.read()
frame = crop_frame(frame)
video_frame = cv2.resize(frame, (250, 250))
cv2.imshow('PRESS "F" TO MARK AS FAIL',video_frame)
k =cv2.waitKey(20)
FAILED_PRINT = False
if(k==ord("f")):
FAILED_PRINT = True
if(count % frame_save_skip == 0):
if FAILED_PRINT:
cv2.imwrite(dataset_PATH+"/1/"+Images_Name+str(int(count/frame_save_skip))+".png", frame)
sys.stdout.write("\033[K")
print("SAVE AS FAILED PRINT!!! : "+str(int(count/frame_save_skip)), end='\r')
else:
cv2.imwrite(dataset_PATH+"/0/"+Images_Name+str(int(count/frame_save_skip))+".png", frame)
sys.stdout.write("\033[K")
print("SAVE AS GOOD PRINT!!! : "+str(int(count/30)), end='\r')
count +=1
if(k == ord('q')):
loop_var = False
sys.stdout.write("\033[K")
print("MAKE DATASET FOR VIDEO "+ video_PATH + " DONE ! "+"MARKED "+str(int(count/frame_save_skip))+" FRAMES !", end='\r')