-
Notifications
You must be signed in to change notification settings - Fork 117
/
main_simple_lib.py
316 lines (261 loc) · 11.9 KB
/
main_simple_lib.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
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
# General imports and variables, as well as config
import ast
import math
import sys
import time
import requests
import torch.multiprocessing as mp
from joblib import Memory
from rich.console import Console
from rich.live import Live
from rich.padding import Padding
from rich.pretty import pprint
from rich.prompt import Prompt
from rich.syntax import Syntax
from rich import print
from rich.markup import escape as rich_escape
from IPython.display import update_display, clear_output, display
from PIL import Image
import matplotlib.pyplot as plt
from configs import config
from utils import show_single_image
from IPython.display import update_display, clear_output
from IPython.core.display import HTML
cache = Memory('cache/' if config.use_cache else None, verbose=0)
mp.set_start_method('spawn', force=True)
from vision_processes import forward, finish_all_consumers # This import loads all the models. May take a while
from image_patch import *
from video_segment import *
from datasets.dataset import MyDataset
console = Console(highlight=False, force_terminal=False)
time_wait_between_lines = 0.5
def inject_saver(code, show_intermediate_steps, syntax=None, time_wait_between_lines=None, console=None):
injected_function_name = 'show_all'
if injected_function_name in code:
return code
code = code.split("\n")
newcode = []
for n, codeline in enumerate(code):
codeline, indent = split_codeline_and_indent_level(codeline)
if codeline.startswith('#') or codeline == '': # this will cause issues if you have lots of comment lines
continue
if '#' in codeline:
codeline = codeline.split('#')[0]
thing_to_show, code_type = get_thing_to_show_codetype(codeline)
if code_type in ('assign', 'append', 'if', 'return', 'for', 'sort', 'add'):
if '\'' in codeline:
codeline.replace('\'', '\\\'')
if show_intermediate_steps:
escape_thing = lambda x: x.replace("'", "\\'")
injection_string_format = \
lambda \
thing: f"{indent}{injected_function_name}(lineno={n},value=({thing}),valuename='{escape_thing(thing)}'," \
f"fig=my_fig,console_in=console,time_wait_between_lines=time_wait_between_lines); " \
f"CodexAtLine({n},syntax=syntax,time_wait_between_lines=time_wait_between_lines)"
else:
injection_string_format = lambda thing: f"{indent}CodexAtLine({n},syntax=syntax," \
f"time_wait_between_lines=time_wait_between_lines)"
extension_list = []
if isinstance(thing_to_show, list):
injection_string_list = [injection_string_format(f"{thing}") for thing in thing_to_show]
extension_list.extend(injection_string_list)
elif code_type == 'for':
injection_string = injection_string_format(f"{thing_to_show}")
injection_string = " " * 4 + injection_string
extension_list.append(injection_string)
else:
extension_list.append(injection_string_format(f"{thing_to_show}"))
if code_type in ('if', 'return'):
extension_list = extension_list + [f"{indent}{codeline}"]
else:
extension_list = [f"{indent}{codeline}"] + extension_list
newcode.extend(extension_list)
elif code_type == 'elif_else':
newcode.append(f"{indent}{codeline}")
else:
newcode.append(f"{indent}{codeline}")
return "\n".join(newcode)
def get_thing_to_show_codetype(codeline):
# can output either a list of things to show, or a single thing to show
things_to_show = []
if codeline.startswith("if"):
condition, rest = codeline[3:].split(":", 1)
codeline = f"if {condition}:{rest}"
code_type = "if"
operators = ['==', '!=', '>=', '<=', '>', '<']
things_to_show = []
for op in operators:
if op in condition:
things_to_show = [x.strip() for x in condition.split(op)]
# print(things_to_show)
break
# things_to_show.append(thing_to_show)
thing_to_show = things_to_show + [condition.strip()]
elif codeline.startswith("for"):
code_type = 'for'
thing_to_show = codeline.split("for ")[1].split(" in ")[0]
elif codeline.startswith("return"):
thing_to_show = codeline.split("return ")[1]
code_type = 'return'
elif ' = ' in codeline:
code_type = 'assign'
thing_to_show = codeline.split(' = ')[0]
elif ' += ' in codeline:
code_type = 'assign'
thing_to_show = codeline.split(' += ')[0]
elif ' -= ' in codeline:
code_type = 'assign'
thing_to_show = codeline.split(' -= ')[0]
elif ' *= ' in codeline:
code_type = 'assign'
thing_to_show = codeline.split(' *= ')[0]
elif ' /= ' in codeline:
code_type = 'assign'
thing_to_show = codeline.split(' /= ')[0]
elif '.append(' in codeline:
code_type = 'append'
thing_to_show = codeline.split('.append(')[0] + '[-1]'
elif '.add(' in codeline:
code_type = 'add'
thing_to_show = codeline.split('.add(')[0]
elif '.sort(' in codeline:
code_type = 'sort'
thing_to_show = codeline.split('.sort(')[0]
elif codeline.startswith("elif") or codeline.startswith("else"):
thing_to_show = None
code_type = 'elif_else'
else:
thing_to_show = None
code_type = 'other'
if isinstance(thing_to_show, list):
thing_to_show = [thing if not (thing.strip().startswith("'") and thing.strip().endswith("'"))
else thing.replace("'", '"') for thing in thing_to_show if thing is not None]
elif isinstance(thing_to_show, str):
thing_to_show = thing_to_show if not (thing_to_show.strip().startswith("'") and
thing_to_show.strip().endswith("'")) else thing_to_show.replace("'", '"')
return thing_to_show, code_type
def split_codeline_and_indent_level(codeline):
origlen = len(codeline)
codeline = codeline.lstrip()
indent = origlen - len(codeline)
indent = " " * indent
return codeline, indent
def show_one_image(image, ax):
if isinstance(image, torch.Tensor):
image = image.detach().cpu()
if image.dtype == torch.float32:
image = image.clamp(0, 1)
image = image.squeeze(0).permute(1, 2, 0)
ax.imshow(image)
def CodexAtLine(lineno, syntax, time_wait_between_lines=1.):
syntax._stylized_ranges = []
syntax.stylize_range('on red', (lineno + 1, 0), (lineno + 1, 80))
time.sleep(time_wait_between_lines)
def show_all(lineno, value, valuename, fig=None, usefig=True, disp=True, console_in=None, time_wait_between_lines=None,
lastlineno=[-1]):
time.sleep(0.1) # to avoid race condition!
if console_in is None:
console_in = console
thing_to_show = value
if lineno is not None and lineno != lastlineno[0]:
console_in.rule(f"[bold]Line {lineno}[/bold]", style="chartreuse2")
lastlineno[0] = lineno # ugly hack
if usefig:
plt.clf()
ax = fig.add_axes([0, 0, 1, 1])
ax.set_xticks([])
ax.set_yticks([])
if isinstance(thing_to_show, Image.Image):
if valuename:
console_in.print(f'{rich_escape(valuename)} = ')
show_one_image(thing_to_show, ax)
elif str(type(thing_to_show)) == "<class 'image_patch.ImagePatch'>":
if valuename:
console_in.print(f'{rich_escape(valuename)} = ')
show_one_image(thing_to_show.cropped_image, ax)
elif isinstance(thing_to_show, list) or isinstance(thing_to_show, tuple):
if len(thing_to_show) > 0:
for i, thing in enumerate(thing_to_show):
disp_ = disp or i < len(thing_to_show) - 1
show_all(None, thing, f"{rich_escape(valuename)}[{i}]", fig=fig, disp=disp_, usefig=usefig)
return
else:
console_in.print(f"{rich_escape(valuename)} is empty")
elif isinstance(thing_to_show, dict):
if len(thing_to_show) > 0:
for i, (thing_k, thing_v) in enumerate(thing_to_show.items()):
disp_ = disp or i < len(thing_to_show) - 1
show_all(None, thing_v, f"{rich_escape(valuename)}['{thing_k}']", fig=fig, disp=disp_, usefig=usefig)
return
else:
console_in.print(f"{rich_escape(valuename)} is empty")
else:
console_in.print(f"{rich_escape(valuename)} = {thing_to_show}")
if time_wait_between_lines is not None:
time.sleep(time_wait_between_lines / 2)
return
# display small
if usefig:
fig.set_size_inches(2, 2)
if disp:
display(fig)
def load_image(path):
if path.startswith("http://") or path.startswith("https://"):
image = Image.open(requests.get(path, stream=True).raw).convert('RGB')
image = transforms.ToTensor()(image)
else:
image = Image.open(path)
image = transforms.ToTensor()(image)
return image
def get_code(query):
model_name_codex = 'codellama' if config.codex.model == 'codellama' else 'codex'
code = forward(model_name_codex, prompt=query, input_type="image")
if config.codex.model not in ('gpt-3.5-turbo', 'gpt-4'):
code = f'def execute_command(image, my_fig, time_wait_between_lines, syntax):' + code # chat models give execute_command due to system behaviour
code_for_syntax = code.replace("(image, my_fig, time_wait_between_lines, syntax)", "(image)")
syntax_1 = Syntax(code_for_syntax, "python", theme="monokai", line_numbers=True, start_line=0)
console.print(syntax_1)
code = ast.unparse(ast.parse(code))
code_for_syntax_2 = code.replace("(image, my_fig, time_wait_between_lines, syntax)", "(image)")
syntax_2 = Syntax(code_for_syntax_2, "python", theme="monokai", line_numbers=True, start_line=0)
return code, syntax_2
def execute_code(code, im, show_intermediate_steps=True):
code, syntax = code
code_line = inject_saver(code, show_intermediate_steps, syntax, time_wait_between_lines, console)
display(HTML("<style>.output_wrapper, .output {height:auto !important; max-height:1000000px;}</style>"))
with Live(Padding(syntax, 1), refresh_per_second=10, console=console, auto_refresh=True) as live:
my_fig = plt.figure(figsize=(4, 4))
try:
exec(compile(code_line, 'Codex', 'exec'), globals())
result = execute_command(im, my_fig, time_wait_between_lines, syntax) # The code is created in the exec()
except Exception as e:
print(f"Encountered error {e} when trying to run with visualizations. Trying from scratch.")
exec(compile(code, 'Codex', 'exec'), globals())
result = execute_command(im, my_fig, time_wait_between_lines, syntax) # The code is created in the exec()
plt.close(my_fig)
def is_not_fig(x):
if x is None:
return True
elif isinstance(x, str):
return True
elif isinstance(x, float):
return True
elif isinstance(x, int):
return True
elif isinstance(x, list) or isinstance(x, tuple):
return all([is_not_fig(xx) for xx in x])
elif isinstance(x, dict):
return all([is_not_fig(xx) for xx in x.values()])
return False
f = None
usefig = False
if not is_not_fig(result):
f = plt.figure(figsize=(4, 4))
usefig = True
console.rule(f"[bold]Final Result[/bold]", style="chartreuse2")
show_all(None, result, 'Result', fig=f, usefig=usefig, disp=False, console_in=console, time_wait_between_lines=0)
def show_single_image(im):
im = Image.fromarray((im.detach().cpu().numpy().transpose(1, 2, 0) * 255).astype("uint8"))
im.copy()
im.thumbnail((400, 400))
display(im)