-
Notifications
You must be signed in to change notification settings - Fork 5
/
Copy pathmain.py
196 lines (157 loc) · 6.64 KB
/
main.py
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
"""
Project Name: ARJewelBox
Author: Asutosh Pati (https://in.linkedin.com/in/asutoshpati)
Date of Creation: 21 Jul 2020
Purpose: Demonstration of AR filters using python
Description:
AR Jewel Box is a virtual jewellery trial application where a person can
trial a new jewellery with out physically handling it. This program
showcase how AI can be used to implement AR. And how AR can be
implemented using python.
ADVISORY:
As this program is created by Asutosh Pati for educational purpose
so use of this project with out authors consent is strictly
prohibited. This project can only be used for teaching purpose
with authors consent and other uses of this program is completely
banned.
Version:
V 0.0.1: Released with beta version
"""
################################ Libraries ################################
import ctypes
import os.path
import uuid
import cv2
import loader
############################# Global variables ############################
# Uncomment the source you want to select;
# for file & rtsp don't forget to change the path in assets\configs\settings.json
# VID_SOURCE = "USB_CAM"
VID_SOURCE = "FILE"
# VID_SOURCE = "RTSP"
# Change to True if you need to flip image
NEED_FLIP = False
# Default width & height of the images
WIDTH = 720
HEIGHT = 640
# Store the current jewellery information
jewellery = None
impose = None
mx = my = None
dw = dh = None
############################# Helper functions ############################
def update_impose():
"""
Update the new jewellery info when previous or next button is pressed.
Returns:
None
"""
global impose, mx, my, dw, dh
impose = jewellery["path"]
mx, my = jewellery["x"], jewellery["y"]
dw, dh = jewellery["dw"], jewellery["dh"]
def save_picture(curr_frame):
"""
Save the current frame as JPEG image.
Args:
curr_frame (np.array): Current frame when user wanted to save it
Returns:
None
"""
# Create the output folder if not exists
if not os.path.exists(loader.get_absolute_path("captures", True)):
os.makedirs(loader.get_absolute_path("captures", True))
filename = r"captures\Capture_"+str(uuid.uuid4())[:6]+".jpg"
filename = loader.get_absolute_path(filename, True)
# Get a unique filename
while os.path.exists(filename):
filename = r"captures\Capture_" + str(uuid.uuid4())[:6]+".jpg"
filename = loader.get_absolute_path(filename, True)
cv2.imwrite(filename, curr_frame)
############################### Main Program ##############################
if __name__ == "__main__":
# Set window icon during .exe execution
MY_APP_ID = "asutosh.arjbx.v0.0.1" # arbitrary string
ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(MY_APP_ID)
# Get the source from settings configuration file
# print(loader.get_source(VID_SOURCE))
cam = cv2.VideoCapture(loader.get_source(VID_SOURCE))
cv2.imshow("preview", loader.generate_loading_screen(HEIGHT, WIDTH))
# Load the required files into memory
cascade, jewelleries = loader.load_files()
jewel_key_list = list(jewelleries)
# Set the first jewellery as current jewellery in use
curr_jewel_index = 0
jewellery = jewelleries[jewel_key_list[curr_jewel_index]]
impose = jewellery["path"]
mx, my = jewellery["x"], jewellery["y"]
dw, dh = jewellery["dw"], jewellery["dh"]
# Start reading the frames from source
# alpha = 0.9 # 0.0 to 1.0
check, frame = cam.read()
cv2.waitKey(2000)
if check:
# If able to read the source; start capturing frames
while check:
check, frame = cam.read()
if not check:
continue
# If you need to flip the image horizontally;
# just change the NEED_FLIP to True above in global variable section
if NEED_FLIP:
frame = cv2.flip(frame, 1)
# Resize the frame
frame = cv2.resize(frame, (WIDTH, HEIGHT))
# cv2.imshow('preview', frame)
# Detect available faces in frame
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
# you can change the scaleFactor & minNeighbors if required
faces = cascade.detectMultiScale(gray, scaleFactor=1.8, minNeighbors=3)
# for x, y, w, h in faces:
# frame = cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 255, 0), 2)
# Show the filter, if only 1 face is found otherwise it may create inconsistency
# You can try with multiple faces... here
if len(faces) == 1:
# Extract face location
x, y, w, h = faces[0]
# Adjust the jewelery to the face's placement
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2BGRA)
fw, fh = int(w * dw), int(w * dh)
new_impose = cv2.resize(impose, (fw, fh))
# Super impose the jewellery over captured frame
iw, ih, c = new_impose.shape
for i in range(0, iw):
for j in range(0, ih):
if new_impose[i, j][3] != 0:
if y + i + h + my < HEIGHT and x + j + mx < WIDTH:
# print(y + i + h + my, x + j + mx)
frame[y + i + h + my, x + j + mx] = new_impose[i, j]
# Display the final image
cv2.imshow("preview", frame)
# Set ESC key to quit the program
k = cv2.waitKey(15)
if k == 27:
break
# Goto next jewellery when N is pressed
elif k == ord("n") or k == ord("N"):
if curr_jewel_index + 1 != len(jewel_key_list):
curr_jewel_index += 1
else:
curr_jewel_index = 0
jewellery = jewelleries[jewel_key_list[curr_jewel_index]]
update_impose()
# Goto previous jewellery when P is pressed
elif k == ord("p") or k == ord("P"):
if curr_jewel_index - 1 < 0:
curr_jewel_index = len(jewel_key_list) - 1
else:
curr_jewel_index -= 1
jewellery = jewelleries[jewel_key_list[curr_jewel_index]]
update_impose()
# Save the picture when S is pressed
elif k == ord("s") or k == ord("S"):
save_picture(frame)
else:
print("Error: Can't read camera")
cv2.destroyAllWindows()
cam.release()