-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
155 lines (128 loc) · 4.2 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
import cv2
import json
import sys
from image_processing import process_frame, get_item_name, has_multiple_variants
import controller
from item import Item
def main(port, init=True):
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('Error opening camera')
sys.exit(1)
# Init controller
ctrl = controller.Controller(port)
ctrl.send_cmd()
if init:
# Make the switch recognize this controller
ctrl.press_button(controller.BTN_A)
ctrl.p_wait(1)
# Exit Change Grip/Order screen
ctrl.press_button(controller.BTN_A)
ctrl.p_wait(1)
# Exit Controllers screen
ctrl.press_button(controller.BTN_B)
ctrl.p_wait(1)
# Resume AC
ctrl.press_button(controller.BTN_HOME)
ctrl.p_wait(1)
# We should be on the "Welcome to Nook Shopping!" screen
# Make sure the cursor is in a known place
ctrl.press_button(controller.DPAD_D)
ctrl.press_button(controller.DPAD_L)
ctrl.press_button(controller.DPAD_L)
ctrl.press_button(controller.DPAD_L)
# Open catalog
ctrl.press_button(controller.BTN_A)
ctrl.p_wait(1)
# Get furniture
furniture = process_screen(ctrl, capture)
with open('furniture.json', 'w') as f:
json.dump(furniture, f, default=lambda o: o.__dict__)
# Get clothing
ctrl.press_button(controller.BTN_B)
ctrl.p_wait(1)
ctrl.press_button(controller.DPAD_R)
ctrl.press_button(controller.BTN_A)
clothing = process_screen(ctrl, capture)
with open('clothing.json', 'w') as f:
json.dump(clothing, f, default=lambda o: o.__dict__)
# Get wallpaper
ctrl.press_button(controller.BTN_B)
ctrl.p_wait(1)
ctrl.press_button(controller.DPAD_R)
ctrl.press_button(controller.BTN_A)
ctrl.p_wait(1)
wallpaper = process_screen(ctrl, capture)
with open('wallpaper.json', 'w') as f:
json.dump(wallpaper, f, default=lambda o: o.__dict__)
capture.release()
def process_screen(ctrl, capture):
# Scroll all the way up
first_item = None
ctrl.send_cmd(controller.RSTICK_U)
ctrl.p_wait(3)
while True:
_, frame = capture.read()
item = get_item_name(frame, 0)
if first_item == item:
break
first_item = item
ctrl.press_button(controller.RSTICK_U)
ctrl.p_wait(0.2)
ctrl.press_button(controller.RSTICK_U)
ctrl.press_button(controller.RSTICK_U)
items = {}
while True:
res = process_item(ctrl, capture, items)
if res is None:
break
ctrl.press_button(controller.DPAD_D)
print(res)
items[res.name] = res
return list(items.values())
def process_item(ctrl, capture, processed_items):
# ctrl.p_wait(0.5)
_, frame = capture.read()
slot = len(processed_items)
if slot > 7:
slot = 7
item_name, has_variants, variant_name = process_frame(frame, slot)
if item_name in processed_items:
return None
item = Item(item_name)
if has_variants:
while not item.has_variant(variant_name):
item.add_variant(variant_name)
ctrl.press_button(controller.BTN_X)
# ctrl.p_wait(0.5)
_, frame = capture.read()
_, _, variant_name = process_frame(frame, only_get_variant=True)
elif variant_name is not None:
item.add_variant(variant_name)
return item
def reset_controller(ctrl=None):
if ctrl is None:
ctrl = controller.Controller('COM6')
ctrl.press_button(controller.BTN_B)
ctrl.press_button(controller.BTN_HOME)
ctrl.p_wait(1)
ctrl.press_button(controller.DPAD_D)
ctrl.press_button(controller.DPAD_R)
ctrl.press_button(controller.DPAD_R)
ctrl.press_button(controller.DPAD_R)
ctrl.press_button(controller.BTN_A)
ctrl.p_wait(1)
ctrl.press_button(controller.BTN_A)
return ctrl
def screen_capture():
capture = cv2.VideoCapture(0)
if not capture.isOpened():
print('Error opening camera')
sys.exit(1)
_, frame = capture.read()
cv2.imwrite('frame.png', frame)
capture.release()
sys.exit()
if __name__ == '__main__':
# main(sys.argv[1])
main('COM6', True)