11
\$\begingroup\$

I've made simple app that is able to take screenshot each -t seconds and OCR it to -o file. It's also able to take -path and OCR all images to -o file.

It works pretty well thanks to tesseract ocr. Unfortunately my code seems much more complicated than it really is. Can you tell me what I'm doing wrong? I enjoy programming, but reading my own code is a nightmare. Should I split it somehow into separate files? How can I improve as a programmer?


from PIL import ImageGrab, Image
from datetime import datetime
from time import sleep

import argparse
import os
from sys import exit

from pyautogui import position
import cv2
import numpy as np
import pytesseract

class screenshoter:
    a,b,c,d = 0, 0, 0, 0
    path = ""
    last_image_sum = 0
    output_file = ""
    ocr_path = ""
    lang = ""

    def __init__(self, output_file, ocr_path, path, lang='pol', a=0, b=0, c=0, d=0 ):
        if c == 0 or d == 0: #screenshot mode
            input("Set mouse cursor on upper left corner and press enter")
            self.a, self.b = position()
            input("OK. Now set mouse cursor on botton right corner and press enter")
            self.c, self.d = position()
        else:
            self.a = a
            self.b = b
            self.c = c
            self.d = d
        self.output_file=output_file
        self.ocr_path=ocr_path
        self.path=path
        self.lang=lang

    def _get(self):
        filename=str(datetime.now().strftime("%d_%m_%H_%M_%S"))
        im = ImageGrab.grab(bbox=(self.a,self.b,self.c,self.d))
        file_with_path = self.path + filename + ".png"
        im.save(file_with_path)
        return file_with_path

    def read(self):
        with open(self.output_file, "a",encoding='utf8') as f:
            filename=str(datetime.now().strftime("%d_%m_%H_%M_%S"))+'tmp.png'
            screenshot = self._get()
            image = cv2.imread(screenshot)
            if np.sum(image) != self.last_image_sum: 
                self.last_image_sum = np.sum(image)
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
                cv2.imwrite(filename, gray)
                pytesseract.pytesseract.tesseract_cmd = self.ocr_path
                text = pytesseract.image_to_string(Image.open(filename), lang=self.lang)
                #this line removesc screenshots, uncomment if you only need text (if ocr fails you will regred)
                #os.remove(screenshot)
                #this line removes files ocr is working on, grayscale and cropped
                os.remove(filename)
                f.write(text)
                print ("Text saved")
            else:
                print ("Skipping, image repeated")

    def read_files(self):
        #read files from dir
        files = [f for f in os.listdir(self.path) if os.path.isfile(os.path.join(self.path, f))]

        with open(self.output_file, "a", encoding='utf8') as f:
            for image in files:
                image = cv2.imread(self.path+image)
                image = image[self.b:self.d, self.a:self.c]
                filename=str(datetime.now().strftime("%d_%m_%H_%M_%S"))+'tmp.png'

                if np.sum(image) != self.last_image_sum: 
                    self.last_image_sum = np.sum(image)
                    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                    gray = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
                    cv2.imwrite(filename, gray)
                    pytesseract.pytesseract.tesseract_cmd = self.ocr_path
                    text = pytesseract.image_to_string(Image.open(filename), lang=self.lang)
                    f.write(text)
                    #this line removes files ocr is working on, grayscale and cropped
                    os.remove(filename)
                    print ("Text saved")
                else:
                    print ("Skipping, image repeated")



if __name__ == "__main__":
    #deal with args
    ap = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, epilog='''\n
        This software has been made to automate taking screenshot and OCRing to text file.
        It is using Tesseract OCR v5.0.0
        Default is screenshot mode with 30s interval.
        To run in files use -m 1 and -c (corners).
        ''')
    ap.add_argument('-m', '--mode', choices=[0, 1], default=0, type=int, help="0 for screenshot, 1 for files. Default is screenshot mode.")
    ap.add_argument('-p', '--path', help="files path, default is .\\images\\")
    ap.add_argument('-c', '--corners', help="Img corners required with files mode eg. -c '0 0 1920 1080' ")
    ap.add_argument('-t', '--time', type=int, help="Screenshot interval, default is 30s")
    ap.add_argument('-o', '--output', help="output file, default is .\\text.txt")
    ap.add_argument('-l', '--lang', help="language, for codes check tesseract docs, default is polish")
    ap.add_argument('-tp', '--tensorpath', help="path to tesseract.exe, default is C:\\Program Files\\Tesseract-OCR\\tesseract.exe")
    args = vars(ap.parse_args())

    #default settings
    output_file="text.txt"
    path='.\\images\\'
    ocr_path="C:\\Program Files\\Tesseract-OCR\\tesseract.exe"
    interval = 30
    lang = 'pol'
    a,b,c,d=0,0,0,0

    ocr = args['tensorpath']
    pth = args['path']
    outp = args['output']
    intrv = args["time"]
    l = args['lang']

    if ocr != None:
        ocr_path = ocr
    if pth != None:
        path = pth
    if outp != None:
        output_file = outp
    if intrv != None:
        interval = intrv
    if l != None:
        lang = l

    #run app
    def run(mode, path, ocr_path, output_file, a,b,c,d, lang):
        if mode:
            app = screenshoter(path=path, ocr_path=ocr_path, output_file=output_file)
            while True:
                app.read()
                sleep(interval)
        else:
            try:
                corners = args["corners"]
                print (corners)
                a,b,c,d = [int(i) for i in corners.split(" ")]
                if c == 0 or d == 0:
                    Exception("-c is required for file mode")
                else:
                    app = screenshoter(a=a, b=b, c=c, d=d, path=path, ocr_path=ocr_path, output_file=output_file)
                    app.read_files()
                    exit(0)
            except Exception as e:
                print (e)
                print ("-c argument is required when parsing from files! eg. -c '0 0 1920 1080'")
                exit(0)

    #mode = True if args['mode'] == None else False
    mode = not args['mode']
    run(mode, path, ocr_path, output_file, a,b,c,d, lang)
\$\endgroup\$
1
  • 1
    \$\begingroup\$ I think the first thing you're doing wrong is to think about screenshots. Is this about OCR from any graphic image, or not? \$\endgroup\$ – Robbie Goodwin May 25 '20 at 17:37
13
\$\begingroup\$

Standard case

Class names like class screenshoter should be capitalized, i.e. class ScreenShoter. The standard Python code style (which includes recommendations such as naming conventions) is documented in the PEP 8 - Style Guide for Python Code.

Static variables

These:

a,b,c,d = 0, 0, 0, 0
path = ""
last_image_sum = 0
output_file = ""
ocr_path = ""
lang = ""

do not do what you think they do. They're effectively statics, and in this case I think they can all be deleted. Whereas you have some defaults here that have not been picked up in your constructor function signature, it doesn't make sense for them to be optional arguments, so your constructor is fine.

Coordinate names

Consider renaming a,b,c,d to x0,x1,y0,y1. It will be easier to understand for other programmers.

Type hints

def __init__(self, output_file, ocr_path, path, lang='pol', a=0, b=0, c=0, d=0 ):

can be, at a guess,

def __init__(self, output_file: str, ocr_path: str, path: str, lang: str='pol', a: float=0, b: float=0, c: float=0, d: float=0):

This helps in a number of ways, including making it much clearer for callers of your constructor.

Redundant str

Do not call str on this:

str(datetime.now().strftime("%d_%m_%H_%M_%S"))

since it's already a string.

pathlib

file_with_path = self.path + filename + ".png"

can be

from pathlib import Path
...
self.path = Path(path)
...
file_with_path = (self.path / filename).with_suffix('.png')

Inline date formatting

str(datetime.now().strftime("%d_%m_%H_%M_%S"))+'tmp.png'

can be

f'{datetime.now():%d_%m_%H_%M_%S}_tmp.png'

Caching

Cache np.sum(image) in a local variable since you use it twice.

Main guard

def run should be moved away from your main guard. It's already in "global scope" but would not be available to unit tests. The other code after your main guard and before run should be moved into its own function.

\$\endgroup\$
10
\$\begingroup\$

The one thing that is immediately apparent to me (before scrolling the code window) is the naming of some variables: a, b, c, d. Why not use more meaningful names ? You're not being billed by length of variable name. If they are coordinates just call them topleft, topright etc.

Ditto with variables names like:

pth = args['path']
outp = args['output']
intrv = args["time"]

Names should be more explicit and not abbreviated needlessly.


I would put all constants at the top of the code.

This should be a constant too: strftime("%d_%m_%H_%M_%S"). Define the desired date format only once.


You are assigning some variables in your code twice eg:

a,b,c,d=0,0,0,0

It is good that you are using argparse but you can do more.

First of all, some variable assignments are unnecessary:

if intrv != None:
    interval = intrv

You can provide a default value for missing arguments. Just use the default option in add_argument. Don't use two variables when one is sufficient.

Since you are expecting 4 numbers for the corners argument here is how to integrate the requirement in argparse: use the nargs option.

import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--corners', dest="corners",
type=int, required=True, nargs=4,
help="Img corners required with files mode eg. -c '0 0 1920 1080' ")

args = parser.parse_args()

# show the values
print(f"corners: {args.corners}")

With this you know that only integers will be accepted. But they could be too large. So you can add one more option like: choices=range(0,1921). With this you will restrict the range of allowed numbers to 0-1920 (as an example). The downside is that Python will output all possibilities if at least one argument does not fulfill the condition. The error message will fill up your screen and this is not pretty.

To avoid this downside, we can write a custom function to validate the argument and override the error message:

import argparse

def check_corners(corner):
    if not corner.isdigit():
        raise argparse.ArgumentTypeError("Corner must be a number")
    if int(corner) not in range(0,1921):
        raise argparse.ArgumentTypeError("Corner must be a number in the range [0-1920]")
    return corner


parser = argparse.ArgumentParser()
parser.add_argument('-c', '--corners', dest="corners",
type=check_corners, required=True, nargs=4,
help="Img corners required with files mode eg. -c '0 0 1920 1080' ")

args = parser.parse_args()

# show the values
print(f"corners: {args.corners}")

There is one caveat though. By default, Python assumes arguments are of type str and that's what happens in this function. So before using range we have to make sure that corner is a number and cast it to int. Otherwise in range will not work.
Or you could use a regex.

Test:

python3 corners.py -c 0 0 0 1921
usage: corners.py [-h] -c CORNERS CORNERS CORNERS CORNERS
corners.py: error: argument -c/--corners: Corner must be a number in the range [0-1920]

With this, you can remove a few lines of your code or at least simplify logic. The argparse module allows you to collect arguments and validate them at the same time. Take full advantage of the features available.


Should I split it somehow into separate files?

Ultimately, keeping the class in a standalone file importable as a packagge would be a good idea.


Use the logging module.

You have a few prints here and there. It would be a good idea to use the logging module instead eg logger.debug("Text saved"). Then you can show the comments on screen, record them to a file or both. You can also reduce verbosity any time you want by adjusting the debug level. If you want your code to become silent you have to comment out all those prints across your module.

Since the code can be automated and is probably unattended to some extent, it is important to keep a log of activity, not just on screen but on file too for later review. Plus, exceptions should be handled and logged too when they occur.


Comments: some functions have no comments at all. Some docstring would be nice. Perhaps add some value samples.

Some functions could be renamed and documented, for example _get. If I'm guessing right it could be take_snapshot or something like that. get is so generic that it is meaningless. The point is that the purpose of each function should be immediately apparent and that's why comments are useful, also for you (in 6 months you'll have to reanalyze your code).


If you work with temp files you could use the tempfile module or write them to /tmp. Then you'll be less concerned with cleanup.

\$\endgroup\$
3
\$\begingroup\$

I echo Reinderien's and Anonymous's comments re your a, b, c, d. I suggest renaming your a and c to left and right, and your b and d to top and bottom. (Two-word names such as topleft are not appropriate because each of these variables holds just one coordinate, not two.)

Then you needn't mandate which corner the user starts from. Let them start from any of the four corners. Then, after you have read x- and y-coordinates of two points given by the user, ensure they're the correct way round:

    if top<bottom:
        top, bottom = bottom, top

(and likewise for left and right). It's better to be generous than needlessly restrictive with ways you enable the user to do things. By contrast, your code needs to be strict about how it does things when that matters, e.g. correctly identifying which y-coor was the top and which one was the bottom.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.