def add_ocrinfo(tree, imgfile):
imgpil = Image.open(imgfile)
(orig_width, orig_height) = (imgpil.width, imgpil.height)
#root_width = tree[min(tree)]['width']
ratio = 1.0 * orig_width / config.width
#imgpil = imgpil.convert("RGB").resize(
# (orig_width * OCR_RATIO, orig_height * OCR_RATIO))
tesapi = PyTessBaseAPI(lang='eng')
tesapi.SetImage(imgpil)
tesapi.SetSourceResolution(config.ocr_resolution)
for nodeid in tree:
node = tree[nodeid]
if node['children'] and node['text'] == '':
node['ocr'] = ''
continue
x = max(node['x'] * ratio - 1, 0)
y = max(node['y'] * ratio - 1, 0)
x2 = min((node['x'] + node['width']) * ratio + 1, orig_width)
y2 = min((node['y'] + node['height']) * ratio + 1, orig_height)
width = int(x2 - x)
height = int(y2 - y)
if width > 3 and height > 3:
#tesapi.SetRectangle(int(x * OCR_RATIO), int(y * OCR_RATIO),
# int(width * OCR_RATIO), int(height * OCR_RATIO))
#print(int(x), int(y), int(width), int(height), orig_width, orig_height)
tesapi.SetRectangle(int(x), int(y), int(width), int(height))
ocr_text = tesapi.GetUTF8Text().strip().replace('\n', ' ')
if ocr_text.strip() == '':
x = min(x + width * 0.05, orig_width)
y = min(y + height * 0.05, orig_height)
width *= 0.9
height *= 0.9
tesapi.SetRectangle(int(x), int(y), int(width), int(height))
ocr_text = tesapi.GetUTF8Text().strip().replace('\n', ' ')
else:
ocr_text = ''
node['ocr'] = ocr_text
class ViewerWindow(Gtk.Window):
def __init__(self, filenames, kind, show, ml):
Gtk.Window.__init__(self)
self.ptx = 0
self.pty = 0
self.focus_id = -1
self.file_idx = 0
self.kind = kind
self.show_hidden = show
self.ml = ml
self.screen_hint = ''
self.in_hint_screen = False
self.colors = {}
self.memory = {}
self.elem_models = {}
self.filenames = filenames
self.tesapi = PyTessBaseAPI(lang='eng')
self.tesapi.SetVariable("tessedit_char_whitelist", WHITELIST)
self.init_ui()
self.load()
def init_ui(self):
self.connect("delete-event", Gtk.main_quit)
darea = Gtk.DrawingArea()
darea.connect("draw", self.on_draw)
darea.connect("motion-notify-event", self.move_over)
darea.connect("button-release-event", self.click_evt)
darea.connect("scroll-event", self.scroll_evt)
darea.connect("key-release-event", self.key_evt)
darea.set_events(Gdk.EventMask.POINTER_MOTION_MASK
| Gdk.EventMask.BUTTON_RELEASE_MASK
| Gdk.EventMask.BUTTON_PRESS_MASK
| Gdk.EventMask.SCROLL_MASK
| Gdk.EventMask.KEY_PRESS_MASK
| Gdk.EventMask.KEY_RELEASE_MASK)
darea.set_can_focus(True)
self.add(darea)
self.show_all()
def load(self, prev=False):
if self.file_idx == len(self.filenames):
Gtk.main_quit()
return
if prev:
self.file_idx -= 2
filename = self.filenames[self.file_idx]
(self.app, self.scr) = util.get_aux_info(filename)
if self.app not in self.memory:
self.memory[self.app] = {}
self.set_title(filename)
self.file_idx += 1
print("Loading %s" % filename)
self.pngfile = os.path.splitext(filename)[0] + '.png'
self.descname = os.path.splitext(filename)[0] + '.%s.txt' % self.kind
starttime = time.time()
self.tree = analyze.load_tree(filename)
hidden.find_hidden_ocr(self.tree)
hidden.mark_children_hidden_ocr(self.tree)
util.print_tree(self.tree, show_hidden=self.show_hidden)
if self.ml:
self.get_ml_rets()
else:
self.load_desc()
endtime = time.time()
print("Load time: %.3fs" % (endtime - starttime))
self.focus_id = -1
self.colors = {}
self.ptx = self.pty = 0
self.img = cairo.ImageSurface.create_from_png(self.pngfile)
print('Image:', self.img.get_width(), self.img.get_height())
root_item_id = min(self.tree)
root_node = self.tree[root_item_id]
print('Root node:', root_node['width'], root_node['height'])
self.scale = 1.0 * self.img.get_width() / config.width
#self.scale = analyze.find_closest(self.scale, analyze.SCALE_RATIOS)
print('Scale:', '%.3f' % self.scale, '->', '%.3f' % self.scale)
self.resize(self.img.get_width(), self.img.get_height())
self.mark_depth(self.tree)
for item_id in self.tree:
color_r = random.random() / 2
color_g = random.random() / 2
color_b = random.random() / 2
self.colors[item_id] = (color_r, color_g, color_b)
imgocr = Image.open(self.pngfile)
self.imgwidth = imgocr.width
self.imgheight = imgocr.height
#imgocr2 = imgocr.convert("RGB").resize(
# (imgocr.width * OCR_RATIO, imgocr.height * OCR_RATIO))
self.tesapi.SetImage(imgocr)
self.tesapi.SetSourceResolution(config.ocr_resolution)
self.dump_memory()
def remember(self, node, desc):
nodeid = node['id']
if not node['id']:
return
if node['id'] in self.memory[self.app]:
if desc != self.memory[self.app][nodeid]:
# multiple!
self.memory[self.app][nodeid] = 'MUL'
else:
self.memory[self.app][node['id']] = desc
def forget(self, node):
if node['id'] in self.memory[self.app]:
del self.memory[self.app][node['id']]
def get_elem_model(self, app):
elem_clas = elements.getmodel("../model/", "../guis/", app,
"../guis-extra/",
config.extra_element_scrs)
self.elem_models[app] = elem_clas
def get_ml_rets(self):
if self.app not in self.elem_models:
self.get_elem_model(self.app)
guess_descs = {}
guess_items = {} # type: Dict[str, List[int]]
guess_score = {}
elem_clas = self.elem_models[self.app]
elem_clas.set_imgfile(self.pngfile)
treeinfo = analyze.collect_treeinfo(self.tree)
for itemid in self.tree:
(guess_element,
score) = elem_clas.classify(self.scr, self.tree, itemid, None,
treeinfo)
if guess_element != 'NONE':
if tags.single(guess_element,
self.scr) and guess_element in guess_items:
old_item = guess_items[guess_element][0]
if guess_score[old_item] < score:
guess_items[guess_element] = [itemid]
guess_score[itemid] = score
del guess_descs[old_item]
guess_descs[itemid] = guess_element
else:
guess_descs[itemid] = guess_element
guess_score[itemid] = score
guess_items[guess_element] = (
guess_items.get(guess_element, []) + [itemid])
for nodeid in guess_descs:
self.tree[nodeid]['label'] = guess_descs[nodeid]
def load_desc(self):
if os.path.exists(self.descname):
with open(self.descname) as inf:
for line in inf.read().split('\n'):
if not line:
continue
(item_id, desc) = line.split(' ', 1)
item_id = int(item_id)
found = False
for nodeid in self.tree:
node = self.tree[nodeid]
if item_id in node['raw']:
if 'label' in node:
node['label'] += ' ' + desc
else:
node['label'] = desc
print(nodeid, '(', item_id, ')', '->', desc)
self.remember(node, desc)
found = True
break
if not found:
print("WARNING: %s (%s) is missing!" % (item_id, desc))
def mark_depth(self, tree):
for item_id in tree:
node = tree[item_id]
if 'depth' in node:
continue
self.mark_depth_node(tree, item_id, 0)
def mark_depth_node(self, tree, node_id, depth):
node = tree[node_id]
node['depth'] = depth
node['descs'] = []
for child in node['children']:
descs = self.mark_depth_node(tree, child, depth + 1)
node['descs'] += descs
return node['descs'] + [node_id]
def get_node_info(self, node):
(x, y, width, height, depth) = (node['x'], node['y'], node['width'],
node['height'], node['depth'])
x *= self.scale
y *= self.scale
width *= self.scale
height *= self.scale
width = min(width, self.imgwidth)
height = min(height, self.imgheight)
if x < 0:
width += x
x = 0
if y < 0:
height += y
y = 0
return (x, y, width, height, depth)
def find_containing_widget(self, px, py):
max_depth = 0
max_id = -1
for item_id in self.tree:
node = self.tree[item_id]
if self.ignore_node(node):
continue
if self.inside(node, px, py):
if node['depth'] > max_depth:
max_depth = node['depth']
max_id = item_id
return max_id
def inside(self, node, px, py):
(x, y, width, height, depth) = self.get_node_info(node)
return x <= px and x + width >= px and y <= py and y + height >= py
def ignore_node(self, node):
if node['class'].upper() == 'OPTION':
return True
if node.get('visible', '') == 'hidden':
return True
return False
def on_draw(self, wid, ctx):
ctx.select_font_face("Arial", cairo.FONT_SLANT_NORMAL,
cairo.FONT_WEIGHT_BOLD)
ctx.set_source_surface(self.img, 0, 0)
ctx.paint()
ctx.set_font_size(20)
ctx.set_line_width(5)
ctx.set_source_rgb(1.0, 0.0, 0.0)
max_click_id = -1
max_click_depth = 0
max_id = self.find_containing_widget(self.ptx, self.pty)
for item_id in self.tree:
node = self.tree[item_id]
depth = node['depth']
if max_id in node['descs'] and node['click']:
if depth > max_click_depth:
max_click_depth = depth
max_click_id = item_id
for item_id in self.tree:
node = self.tree[item_id]
if self.ignore_node(node):
continue
if item_id == max_id:
region_mode = False
else:
region_mode = True
(x, y, width, height, depth) = self.get_node_info(node)
if not self.inside(node, self.ptx, self.pty):
continue
self.show_widget(ctx, item_id, not region_mode, not region_mode)
if max_click_id != -1 and max_click_id != max_id:
self.show_widget(ctx, max_click_id, False, True)
if self.focus_id >= 0:
self.show_widget(ctx, self.focus_id, True, True, (1, 0, 0))
for itemid in self.tree:
node = self.tree[itemid]
if 'label' in node:
if itemid == self.focus_id:
color = (0, 1, 0)
else:
color = (0, 0, 1)
self.show_widget(ctx, itemid, True, False, (0, 0, 1))
self.show_desc(ctx, node, color)
#s.write_to_png('test.png')
#os.system("%s %s" % (config.picviewer_path, 'test.png'))
#report_time(start_time, "displayed")
def move_sibling(self, to_next):
leaf_list = []
any_list = []
for itemid in self.tree:
node = self.tree[itemid]
if not self.inside(node, self.clickx, self.clicky):
continue
if len(node['children']) == 0:
leaf_list.append(itemid)
any_list.append(itemid)
for i in range(len(leaf_list)):
if leaf_list[i] == self.focus_id:
if to_next:
idx = (i + 1) % len(leaf_list)
else:
idx = (i - 1) % len(leaf_list)
self.focus_id = leaf_list[idx]
return
if len(leaf_list) == 0:
for i in range(len(any_list)):
if any_list[i] == self.focus_id:
if to_next:
idx = (i + 1) % len(any_list)
else:
idx = (i - 1) % len(any_list)
self.focus_id = any_list[idx]
return
self.focus_id = any_list[0]
else:
self.focus_id = leaf_list[0]
def show_widget(self, ctx, item_id, fill, show_text, colors=None):
node = self.tree[item_id]
(x, y, width, height, depth) = self.get_node_info(node)
if colors is None:
color_r = self.colors[item_id][0]
color_g = self.colors[item_id][1]
color_b = self.colors[item_id][2]
else:
(color_r, color_g, color_b) = colors
ctx.rectangle(x, y, width, height)
if fill:
ctx.set_source_rgba(color_r, color_g, color_b, 0.3)
ctx.fill()
else:
ctx.set_source_rgba(color_r, color_g, color_b, 1)
ctx.set_line_width(5)
ctx.stroke()
if show_text:
max_char = int(width / ctx.text_extents("a")[2])
text = str(item_id)
if node['click']:
text = 'C' + text
if node['text']:
text = text + ':' + node['text'][:(max_char - 5)]
elif node['id']:
text += '#' + node['id'][:(max_char - 5)]
self.show_text(ctx, x + width / 2, y + height / 2, text, color_r,
color_g, color_b)
def show_desc(self, ctx, node, color=(0, 0, 1)):
desc = node['label']
(x, y, width, height, depth) = self.get_node_info(node)
self.show_text(ctx, x + width / 2, y + height / 2, desc, color[0],
color[1], color[2])
def show_text(self, ctx, x, y, text, color_r, color_g, color_b):
x_bearing, y_bearing, text_width, text_height = ctx.text_extents(
text)[:4]
ctx.move_to(x - text_width / 2, y + text_height / 2)
ctx.set_source_rgba(1, 1, 1, 1)
ctx.set_line_width(5)
ctx.text_path(text)
ctx.stroke()
ctx.move_to(x - text_width / 2, y + text_height / 2)
ctx.set_source_rgba(color_r, color_g, color_b, 1)
ctx.text_path(text)
ctx.fill()
def move_over(self, widget, evt):
self.ptx = evt.x
self.pty = evt.y
self.queue_draw()
def click_evt(self, widget, evt):
if self.in_hint_screen:
self.process_screen_hint_click(evt)
return
if evt.button == 3:
self.focus_id = -1
else:
self.clickx = evt.x
self.clicky = evt.y
self.focus_id = self.find_containing_widget(evt.x, evt.y)
self.queue_draw()
def scroll_evt(self, widget, evt):
if self.focus_id == -1:
return
scroll_up = evt.direction == Gdk.ScrollDirection.UP
if scroll_up:
self.focus_id = self.find_parent_widget(self.focus_id)
else:
self.focus_id = self.find_child_widget(self.focus_id)
self.queue_draw()
def find_parent_widget(self, wid):
for itemid in self.tree:
node = self.tree[itemid]
if self.ignore_node(node):
continue
if wid in node['children']:
return itemid
return wid
def find_child_widget(self, wid):
for itemid in self.tree[wid]['children']:
node = self.tree[itemid]
if self.ignore_node(node):
continue
if self.inside(node, self.clickx, self.clicky):
return itemid
return wid
def mark_direct(self):
enter = self.get_text('Please enter id_label', 'format: <id> <label>')
if enter is None:
return
if ' ' in enter:
nodeid, label = enter.split(' ')
else:
nodeid = enter
label = ''
nodeid = int(nodeid)
if nodeid not in self.tree:
print('missing node', nodeid)
return
node = self.tree[nodeid]
self.mark_node(node, label)
def mark_focused(self):
if self.focus_id < 0:
return
node = self.tree[self.focus_id]
label = self.get_text(
'Please enter label', 'label for %s: %s (%s) #%s' %
(self.focus_id, node['text'], node['desc'], node['id']))
if label is None:
return
if self.ml:
if label == '':
if 'label' not in self.tree[self.focus_id]:
return
self.generate_negative_hint(self.tree[self.focus_id]['label'])
del self.tree[self.focus_id]['label']
else:
self.generate_hint_for_widget(self.focus_id, label)
self.add_label(node, label)
else:
self.mark_node(node, label)
def generate_hint_for_widget(self, nodeid, label):
return self.generate_hint(label,
locator.get_locator(self.tree, nodeid))
def generate_negative_hint(self, label):
return self.generate_hint(label, 'notexist')
def generate_hint(self, label, hint):
print("@%s.%s %s" % (self.scr, label, hint))
def mark_node(self, node, label):
if label == '':
if 'label' in node:
del node['label']
self.forget(node)
else:
self.add_label(node, label)
self.remember(node, label)
self.save_labels()
def ocr_text(self):
node = self.tree[self.focus_id]
(x, y, width, height, _) = self.get_node_info(node)
print(x, y, width, height)
x = max(x - 1, 0)
y = max(y - 1, 0)
width = min(width + 2, self.imgwidth)
height = min(height + 2, self.imgheight)
#self.tesapi.SetRectangle(x * OCR_RATIO, y * OCR_RATIO,
# width * OCR_RATIO, height * OCR_RATIO)
self.tesapi.SetRectangle(x, y, width, height)
print("OCR ret:", self.tesapi.GetUTF8Text())
x = min(x + width * 0.05, self.imgwidth)
y = min(y + height * 0.05, self.imgheight)
width *= 0.9
height *= 0.9
self.tesapi.SetRectangle(x, y, width, height)
print("OCR ret:", self.tesapi.GetUTF8Text())
def save_region(self):
if self.focus_id == -1:
return
node = self.tree[self.focus_id]
(x, y, width, height, _) = self.get_node_info(node)
x = max(x - 1, 0)
y = max(y - 1, 0)
width = min(width + 2, self.imgwidth)
height = min(height + 2, self.imgheight)
regimg = cairo.ImageSurface(cairo.FORMAT_RGB24, int(width),
int(height))
ctx = cairo.Context(regimg)
ctx.set_source_surface(self.img, -x, -y)
ctx.paint()
regimg.write_to_png("/tmp/region.png")
def dump_memory(self):
for _id in self.memory[self.app]:
print('MEM %s -> %s' % (_id, self.memory[self.app][_id]))
def add_label(self, node, desc):
print('%s -> %s' % (util.describe_node(node, short=True), desc))
node['label'] = desc
def auto_label(self):
for nodeid in self.tree:
node = self.tree[nodeid]
if 'label' not in node and node['id'] in self.memory[self.app]:
if self.memory[self.app][node['id']] != 'MUL':
self.add_label(node, self.memory[self.app][node['id']])
else:
print('skip MUL id: %s' % node['id'])
self.save_labels()
def remove_all(self):
for nodeid in self.tree:
node = self.tree[nodeid]
if 'label' in node:
del node['label']
def process_screen_hint_click(self, evt):
click_id = self.find_containing_widget(evt.x, evt.y)
if click_id == -1:
print('Invalid widget selected')
return
hint = locator.get_locator(self.tree, click_id)
if hint is None:
print('Cannot generate hint for this widget')
return
hint = str(hint)
if evt.button == 3:
# negate
hint = 'not ' + hint
print('Widget hint: "%s"' % hint)
self.add_screen_hint(hint)
def add_screen_hint(self, hint):
if self.screen_hint == '':
self.screen_hint = hint
else:
self.screen_hint += ' && ' + hint
def hint_screen(self):
if not self.in_hint_screen:
label = self.get_text('Please enter screen name',
'screen name like "signin"')
if label is None:
return
self.screen_hint_label = label
self.in_hint_screen = True
self.screen_hint = ''
else:
self.in_hint_screen = False
print("%%%s %s" % (self.screen_hint_label, self.screen_hint))
def key_evt(self, widget, evt):
if evt.keyval == Gdk.KEY_space:
self.mark_focused()
elif evt.keyval == Gdk.KEY_Tab:
self.load()
elif evt.keyval == Gdk.KEY_Left:
self.move_sibling(to_next=True)
elif evt.keyval == Gdk.KEY_Right:
self.move_sibling(to_next=False)
elif evt.keyval == Gdk.KEY_v:
self.ocr_text()
elif evt.keyval == Gdk.KEY_a:
self.auto_label()
elif evt.keyval == Gdk.KEY_p:
self.load(prev=True)
elif evt.keyval == Gdk.KEY_l:
self.mark_direct()
elif evt.keyval == Gdk.KEY_r:
self.remove_all()
elif evt.keyval == Gdk.KEY_s:
self.save_region()
elif evt.keyval == Gdk.KEY_x:
self.hint_screen()
self.queue_draw()
def save_labels(self):
with open(self.descname, 'w') as outf:
for itemid in sorted(self.tree):
node = self.tree[itemid]
if 'label' in node:
outf.write("%s %s\n" % (itemid, node['label']))
def get_text(self, title, prompt):
#base this on a message dialog
dialog = Gtk.MessageDialog(self, 0, Gtk.MessageType.QUESTION,
Gtk.ButtonsType.OK_CANCEL, title)
dialog.format_secondary_text(prompt)
#create the text input field
entry = Gtk.Entry()
#allow the user to press enter to do ok
entry.connect("activate",
lambda entry: dialog.response(Gtk.ResponseType.OK))
#create a horizontal box to pack the entry and a label
hbox = Gtk.HBox()
hbox.pack_start(Gtk.Label("Label:"), False, 5, 5)
hbox.pack_end(entry, True, 0, 0)
#add it and show it
dialog.vbox.pack_end(hbox, True, True, 0)
dialog.show_all()
#go go go
response = dialog.run()
if response == Gtk.ResponseType.OK:
text = entry.get_text()
else:
text = None
dialog.destroy()
return text