class PoolTableChecker: def __init__(self): self.tableFinder = PoolTableFinder() self.imageComparator = ImageComparator() self.imageHandler = ImageHandler() def isTableFree(self, imagePathA, imagePathB): imageA = self.imageHandler.readImage(imagePathA) imageB = self.imageHandler.readImage(imagePathB) if CommandHandler.areBoundsSet(): bounds = CommandHandler.getBounds() croppedA = self.imageHandler.cropAndRotateRectangleInImage( imageA, bounds) croppedB = self.imageHandler.cropAndRotateRectangleInImage( imageB, bounds) else: tableABounds = self.tableFinder.getBoundingRectForTable(imageA) tableBBounds = self.tableFinder.getBoundingRectForTable(imageB) largestBounds = self.imageHandler.getLargestBoundingRect( tableABounds, tableBBounds) croppedA = self.imageHandler.cropImageByBoundingRect( imageA, largestBounds) croppedB = self.imageHandler.cropImageByBoundingRect( imageB, largestBounds) return self.imageComparator.areImagesEqual(croppedA, croppedB)
class PoolTableFinder: def __init__(self): self.imageHandler = ImageHandler() self.rawImagePath = ConfigHandler.getRawImagePath() self.croppedImagePath = ConfigHandler.getCroppedImagePath() def saveCroppedImage(self): Logger.info("Saving cropped image") currentImage = self.imageHandler.readImage(self.rawImagePath) if CommandHandler.areBoundsSet(): bounds = CommandHandler.getBounds() croppedImage = self.imageHandler.cropAndRotateRectangleInImage( currentImage, bounds) else: bounds = self.getBoundingRectForTable(currentImage) croppedImage = self.imageHandler.cropImageByBoundingRect( currentImage, bounds) cv.imwrite(self.croppedImagePath, croppedImage) def getMaskForTable(self, poolImage): self.th_green_low = numpy.array(CommandHandler.getThGreenLow()) self.th_green_high = numpy.array(CommandHandler.getThGreenHigh()) hsvImage = cv.cvtColor(poolImage, cv.COLOR_BGR2HSV) return self.getMask(hsvImage) def getBoundingRectForTable(self, poolImage): mask = self.getMaskForTable(poolImage) contours = self.getContours(mask) return self.getLargestBoundingRect(contours) def getMask(self, hsvImage): mask = cv.inRange(hsvImage, self.th_green_low, self.th_green_high) return mask def getContours(self, mask): contours = cv.findContours(mask.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)[-2] if len(contours) == 0: raise PoolTableException("No contours found") return contours def getLargestBoundingRect(self, contours): area = max(contours, key=cv.contourArea) return cv.boundingRect(area)