pagenet/pageCropAnnotator.py

#!/usr/bin/python

#This takes a list of image files and acts as a tool to mark the crop region of the page

import re
import xml.etree.ElementTree as ET
import os
import sys
from io import StringIO
import cv2



def showControls():
    print(' -----------------------------------------------')
    print('| CONTROLS:                                     |')
    print('| * set new corner (base on loc):   left-click  |')
    print('| * set new seam corner (two-page): middle-click|')
    print('| * confirm corners:                enter       |')
    print('| * mark page as abnormal:          a           |')
    print('| * undo:                           backspace   |')
    print('| * start previous page over:       backspace(+)|')
    #print('| * start current page over:        delete      |')
    print('| * exit:                           esc         |')
    print('|                                               |')
    print(' -----------------------------------------------')



lastDidList=[]
tl=(-1,-1)
tr=(-1,-1)
bl=(-1,-1)
br=(-1,-1)
tm=(-1,-1)
bm=(-1,-1)

image=None
orig=None

abnorm=False

def draw():
    global image,tl,tr,bl,br,tm,bm,abnorm
    if tm[0]>=0 and bm[0]>=0:
        cv2.line(image, tm, bm, (0,255,0), 2)
    if tm[0]>=0 and tl[0]>=0:
        cv2.line(image, tm, tl, (0,255,0), 2)
    if tm[0]>=0 and tr[0]>=0:
        cv2.line(image, tm, tr, (0,255,0), 2)
    if bm[0]>=0 and bl[0]>=0:
        cv2.line(image, bm, bl, (0,255,0), 2)
    if bm[0]>=0 and br[0]>=0:
        cv2.line(image, bm, br, (0,255,0), 2)
    if tl[0]>=0 and tr[0]>=0:
        cv2.line(image, tl, tr, (255,0,0), 2)
    if br[0]>=0 and tr[0]>=0:
        cv2.line(image, br, tr, (255,0,0), 2)
    if br[0]>=0 and bl[0]>=0:
        cv2.line(image, br, bl, (255,0,0), 2)
    if tl[0]>=0 and bl[0]>=0:
        cv2.line(image, tl, bl, (255,0,0), 2)


    if tl[0]>=0:
        image[tl[1],tl[0]]=(0,0,255)
        cv2.circle(image, tl, 2, (0,0,200), 1)
        cv2.circle(image, tl, 5, (0,0,200), 2)
    if tr[0]>=0:
        image[tr[1],tr[0]]=(0,0,255)
        cv2.circle(image, tr, 2, (0,0,200), 1)
        cv2.circle(image, tr, 5, (0,0,200), 2)
    if bl[0]>=0:
        image[bl[1],bl[0]]=(0,0,255)
        cv2.circle(image, bl, 2, (0,0,200), 1)
        cv2.circle(image, bl, 5, (0,0,200), 2)
    if br[0]>=0:
        image[br[1],br[0]]=(0,0,255)
        cv2.circle(image, br, 2, (0,0,200), 1)
        cv2.circle(image, br, 5, (0,0,200), 2)
    if tm[0]>=0:
        image[tm[1],tm[0]]=(0,0,255)
        cv2.circle(image, tm, 2, (0,100,200), 1)
        cv2.circle(image, tm, 5, (0,100,200), 2)
    if bm[0]>=0:
        image[bm[1],bm[0]]=(0,0,255)
        cv2.circle(image, bm, 2, (0,100,200), 1)
        cv2.circle(image, bm, 5, (0,100,200), 2)

    if abnorm:
        cv2.putText(image, 'ABNORMAL', (image.shape[1]/2,image.shape[0]/2), cv2.FONT_HERSHEY_PLAIN, 2, (0,0,255))


    cv2.imshow("image", image)

bimage=None
def clicker(event, x, y, flags, param):
        # grab references to the global variables
        global image,tl,tr,bl,br,tm,bm,lastDidList,orig

        """if event == cv2.EVENT_LBUTTONDOWN:
            if len(segPts)>0:
                #change last boundary
                image=bimage.copy()
                segPts[-1]=x
                ll=max(0,segPts[-1]-1)
                rr=min(image.shape[1], segPts[-1]+1)
                image[:,ll:rr,0] = color[(colorIdx+len(color)-1)%len(color)][0] * image[:,ll:rr,0]
                image[:,ll:rr,1] = color[(colorIdx+len(color)-1)%len(color)][1] * image[:,ll:rr,1]
                image[:,ll:rr,2] = color[(colorIdx+len(color)-1)%len(color)][2] * image[:,ll:rr,2]
                cv2.imshow("image", image)
"""
        if event == cv2.EVENT_LBUTTONDOWN:
                # a new boundary
                if x<image.shape[1]/2 and y<image.shape[0]/2:
                    tl=(x,y)
                    if 0 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(0)
                    lastDidList.append(0)
                if x>image.shape[1]/2 and y<image.shape[0]/2:
                    tr=(x,y)
                    if 1 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(1)
                    lastDidList.append(1)
                if x<image.shape[1]/2 and y>image.shape[0]/2:
                    bl=(x,y)
                    if 2 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(2)
                    lastDidList.append(2)
                if x>image.shape[1]/2 and y>image.shape[0]/2:
                    br=(x,y)
                    if 3 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(3)
                    lastDidList.append(3)
                draw()

        elif event == cv2.EVENT_MBUTTONDOWN:
                # a new boundary
                if y<image.shape[0]/2:
                    tm=(x,y)
                    if 1 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(1)
                    lastDidList.append(4)
                if y>image.shape[0]/2:
                    bm=(x,y)
                    if 3 in lastDidList:
                        image=orig.copy()
                        lastDidList.remove(3)
                    lastDidList.append(5)
                draw()

def segmenter(imDir,imagePath,dispHeight):
    global image,tl,tr,bl,br,tm,bm,lastDidList,orig,abnorm
    print('opening '+imDir+imagePath)
    orig = cv2.imread(imDir+imagePath)
    scale = orig.shape[0]/dispHeight
    orig = cv2.resize(orig,(0,0),None,1.0/scale,1.0/scale)
    #print 'opened'
    assert orig is not None
    redo=True
    while redo: #undo loop
        abnorm=False
        lastDidList=[]
        tl=(-1,-1)
        tr=(-1,-1)
        bl=(-1,-1)
        br=(-1,-1)
        tm=(-1,-1)
        bm=(-1,-1)
        redo=False
        image = orig.copy()
        draw()
        while True:
            # display the imageWork and wait for a keypress
            key = cv2.waitKey(33) & 0xFF #so it is robust on all systems
            #print key
            if key == 13 and tl[0]>=0 and tr[0]>=0 and bl[0]>=0 and br[0]>=0: #enter
                toWrite = imagePath+','+str(int(scale*tl[0]))+','+str(int(scale*tl[1]))+','+str(int(scale*tr[0]))+','+str(int(scale*tr[1]))+','+str(int(scale*br[0]))+','+str(int(scale*br[1]))+','+str(int(scale*bl[0]))+','+str(int(scale*bl[1]))
                if abnorm:
                    if tm[0]>=0 and bm[0]>=0:
                        toWrite += ',ABNORMAL,'+str(int(scale*tm[0]))+','+str(int(scale*tm[1]))+','+str(int(scale*bm[0]))+','+str(int(scale*bm[1]))
                    else:
                        toWrite += ',ABNORMAL'
                else:
                    if tm[0]>=0 and bm[0]>=0:
                        toWrite += ',DOUBLE,'+str(int(scale*tm[0]))+','+str(int(scale*tm[1]))+','+str(int(scale*bm[0]))+','+str(int(scale*bm[1]))
                    else:
                        toWrite += ',SINGLE'
                toWrite+='\n';
                return toWrite, False, False
            elif key == 8: #backspace
                if len(lastDidList)>0:
                    imageWork = orig.copy()
                    lastDid=lastDidList.pop()
                    if lastDid==0:
                        tl=(-1,-1)
                    elif lastDid==1:
                        tr=(-1,-1)
                    elif lastDid==2:
                        bl=(-1,-1)
                    elif lastDid==3:
                        br=(-1,-1)
                    elif lastDid==4:
                        tm=(-1,-1)
                    elif lastDid==5:
                        bm=(-1,-1)
                    image=orig.copy()
                    draw()
                else:
                    return '', True, False
            elif key == 127: #del
                #if len(lastDidList)>0:
                    print('[CLEAR]')
                    redo=True
                    break
                #else:
                #    return '', True, False
            elif key == 27: #esc
                print('esc')
                return '', False, True
                #exit(0)
                #break
            elif key == 97: #'a'
                #return imagePath+',-1,-1,-1,-1,-1,-1,-1,-1,ABNORMAL\n', False, False
                abnorm = not abnorm
                image=orig.copy()
                draw()

    #return newWords, newWordBoxes

if len(sys.argv)<4:
    print('usage: '+sys.argv[0]+' imgDir imgList outAnn.csv [displayHeight]')
    print('output format: imageFile, tlx, tly, trx, try, brx, bry, blx, bly, type (,tmx, tmy, bmx, bmy)')
    exit(0)

inFile = sys.argv[2]
imDir = sys.argv[1]
if imDir[-1]!='/':
    imDir+='/'
outFile = sys.argv[3]
dispHeight=500.0
if len(sys.argv)>4:
    dispHeight=float(sys.argv[4])

cv2.namedWindow("image")
cv2.setMouseCallback("image", clicker)

didCount=0
did=[]
try:
    check = open(outFile,'r')
    did = check.read().splitlines()
    didCount=len(did)
    check.close()
    print('found '+outFile+', appending. Note: this is sychronizing based on count alone, if '+inFile+' hash changed, but sure to align '+outFile)
except IOError:
    print(('making new out:'+outFile))

out = open(outFile,'w')

print(' =============================================== ')
print(' !!! INSTRUCTIONS !!!')
print(' If the page does not contain a single page, or  ')
print(' an open book, mark as abnormal with INSERT (e.g.')
print(' two seperate pages).')
print(' Click on the four corners to include all the    ')
print(' full pages in the image (including two pages if ')
print(' fully present).')
print(' If two pages a full present also mark page seam ')
print(' (middle-click).')
print(' On placing points, prioritize the following to  ')
print(' be included/discluded from the polygons in the  ')
print(' following order:')
print('   1. Including the present page(s) content.')
print('   2. Discluding other pages and background.')
print('   3. Discluding the present page(s) boudary.')
print('   4. Including the present page(s) white area.')
#print ' book). If a corner is torn, click where it ought'
#print ' to be, based on page edges. The page seem on an '
#print ' open book is the page edge.'
print(' Use ESC to exit or the latest page you finished ')
print(' will be lost.')

#i=didCount
i=0
#pageCount=-1
prevSeg=''
seg=''
showControls()
inF = open(inFile,'r')
images = inF.read().splitlines()
end=False
doneOne=False
while i<len(images) and not end:
    if i%10==9:
        showControls()
    print((str(i+1)+' of '+str(len(images))))

    if len(did)>i:
        line = did[i].strip().split(',')
        typ = line[8]
        #print typ
        if typ != '-1':
            out.write(did[i].strip()+'\n')
            i+=1
            continue
        seg, undo, end = segmenter(imDir, images[i],dispHeight)
        out.write(seg)
        seg=''
        i+=1
    else:


        seg, undo, end = segmenter(imDir, images[i],dispHeight)
        if len(seg)>0:
            doneOne=True
        
        if undo and i>0 and doneOne:
            prevSeg=''
            print((str(i)+' of '+str(len(images))))
            prevSeg, undo, end = segmenter(imDir, images[i-1],dispHeight)
        else:
            out.write(prevSeg)
            prevSeg=seg
            seg=''
            i+=1
out.write(prevSeg)
out.write(seg)