Skip to content

Commit

Permalink
Add ScreenshotState class for screenshots comparison purposes
Browse files Browse the repository at this point in the history
  • Loading branch information
Mykola Mokhnach committed Mar 11, 2017
1 parent 48f3a96 commit d6a368c
Show file tree
Hide file tree
Showing 3 changed files with 503 additions and 0 deletions.
1 change: 1 addition & 0 deletions build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ dependencies {
compile 'commons-io:commons-io:2.5'
compile 'org.springframework:spring-context:4.3.5.RELEASE'
compile 'org.aspectj:aspectjweaver:1.8.10'
compile 'nu.pattern:opencv:2.4.9-7'

testCompile 'junit:junit:4.12'
}
Expand Down
326 changes: 326 additions & 0 deletions src/main/java/io/appium/java_client/ScreenshotState.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,326 @@
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package io.appium.java_client;

import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.Size;
import org.opencv.imgproc.Imgproc;

import java.awt.AlphaComposite;
import java.awt.Graphics2D;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.util.Optional;
import java.util.function.Function;
import java.util.function.Supplier;

public class ScreenshotState {
private static final long DEFAULT_INTERVAL_MS = 500;

private Optional<BufferedImage> previousScreenshot = Optional.empty();
private Supplier<BufferedImage> stateProvider;

private long comparisonIntervalMs = DEFAULT_INTERVAL_MS;

/**
* The class constructor accepts single argument, which is
* lambda function, that provides the screenshot of the necessary
* screen area to be verified for similarity.
* This lambda method is NOT called upon class creation.
* One has to invoke {@link #remember()} method in order to call it.
*
* <p>Examples of provider function with Appium driver:
* <code>
* () -&gt; {
* final byte[] srcImage = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);
* return ImageIO.read(new ByteArrayInputStream(srcImage));
* }
* </code>
* or
* <code>
* () -&gt; {
* final byte[] srcImage = ((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES);
* final BufferedImage screenshot = ImageIO.read(new ByteArrayInputStream(srcImage));
* final WebElement element = driver.findElement(locator);
* // Can be simplified in Selenium 3.0+ by using getRect method of WebElement interface
* final Point elementLocation = element.getLocation();
* final Dimension elementSize = element.getSize();
* return screenshot.getSubimage(
* new Rectangle(elementLocation.x, elementLocation.y, elementSize.width, elementSize.height);
* }
* </code>
*
* @param stateProvider lambda function, which returns a screenshot for further comparison
*/
public ScreenshotState(Supplier<BufferedImage> stateProvider) {
this.stateProvider = stateProvider;
}

/**
* Gets the interval value in ms between similarity verification rounds in <em>verify*</em> methods.
*
* @return current interval value in ms
*/
public long getComparisonInterval() {
return comparisonIntervalMs;
}

/**
* Sets the interval between similarity verification rounds in <em>verify*</em> methods.
*
* @param ms interval value in ms. 500 by default
* @return self instance for chaining
*/
public ScreenshotState setComparisonInterval(long ms) {
this.comparisonIntervalMs = ms;
return this;
}

/**
* Call this method to save the initial screenshot state.
* It is mandatory to call before any <em>verify*</em> method is invoked.
*
* @return self instance for chaining
*/
public ScreenshotState remember() {
this.previousScreenshot = Optional.of(stateProvider.get());
return this;
}

/**
* This method allows to pass a custom bitmap for further comparison
* instead of taking one using screenshot provider function. This might
* be useful in some advanced cases.
*
* @param customInitialState valid bitmap
* @return self instance for chaining
*/
public ScreenshotState remember(BufferedImage customInitialState) {
this.previousScreenshot = Optional.of(customInitialState);
return this;
}

public static class ScreenshotComparisonError extends RuntimeException {
private static final long serialVersionUID = -7011854909939194466L;

ScreenshotComparisonError(Throwable reason) {
super(reason);
}

ScreenshotComparisonError(String message) {
super(message);
}
}

public static class ScreenshotComparisonTimeout extends RuntimeException {
private static final long serialVersionUID = 6336247721154252476L;
private double currentScore = Double.NaN;

ScreenshotComparisonTimeout(String message, double currentScore) {
super(message);
this.currentScore = currentScore;
}

public double getCurrentScore() {
return currentScore;
}
}

private ScreenshotState checkState(Function<Double, Boolean> checkerFunc, long timeoutMs) {
return checkState(checkerFunc, timeoutMs, ResizeMode.NO_RESIZE);
}

private ScreenshotState checkState(Function<Double, Boolean> checkerFunc, long timeoutMs, ResizeMode resizeMode) {
final long started = System.currentTimeMillis();
double score;
do {
final BufferedImage currentState = stateProvider.get();
score = getOverlapScore(this.previousScreenshot
.orElseThrow(() -> new ScreenshotComparisonError("Initial screenshot state is not set. "
+ "Nothing to compare")), currentState, resizeMode);
if (checkerFunc.apply(score)) {
return this;
}
try {
Thread.sleep(comparisonIntervalMs);
} catch (InterruptedException e) {
throw new ScreenshotComparisonError(e);
}
}
while (System.currentTimeMillis() - started <= timeoutMs);
throw new ScreenshotComparisonTimeout(
String.format("SScreenshot comparison timed out after %s ms. Actual similarity score: %.5f",
timeoutMs, score), score);
}

/**
* Verifies whether the state of the screenshot provided by stateProvider lambda function
* is changed within the given timeout.
*
* @param timeoutMs timeout value in milliseconds
* @param minScore the value in range (0.0, 1.0)
* @return self instance for chaining
* @throws ScreenshotComparisonTimeout if the calculated score is still
* greater or equal to the given score after timeout happens
* @throws ScreenshotComparisonError if {@link #remember()} method has not been invoked yet
*/
public ScreenshotState verifyChanged(long timeoutMs, double minScore) {
return checkState((x) -> x < minScore, timeoutMs);
}

/**
* Verifies whether the state of the screenshot provided by stateProvider lambda function
* is changed within the given timeout.
*
* @param timeoutMs timeout value in milliseconds
* @param minScore the value in range (0.0, 1.0)
* @param resizeMode one of <em>ResizeMode</em> enum values.
* Set it to a value different from <em>NO_RESIZE</em>
* if the actual screenshot is expected to have different
* dimensions in comparison to the previously remembered one
* @return self instance for chaining
* @throws ScreenshotComparisonTimeout if the calculated score is still
* greater or equal to the given score after timeout happens
* @throws ScreenshotComparisonError if {@link #remember()} method has not been invoked yet
*/
public ScreenshotState verifyChanged(long timeoutMs, double minScore, ResizeMode resizeMode) {
return checkState((x) -> x < minScore, timeoutMs, resizeMode);
}

/**
* Verifies whether the state of the screenshot provided by stateProvider lambda function
* is not changed within the given timeout.
*
* @param timeoutMs timeout value in milliseconds
* @param minScore the value in range (0.0, 1.0)
* @return self instance for chaining
* @throws ScreenshotComparisonTimeout if the calculated score is still
* less than the given score after timeout happens
* @throws ScreenshotComparisonError if {@link #remember()} method has not been invoked yet
*/
public ScreenshotState verifyNotChanged(long timeoutMs, double minScore) {
return checkState((x) -> x >= minScore, timeoutMs);
}

/**
* Verifies whether the state of the screenshot provided by stateProvider lambda function
* is changed within the given timeout.
*
* @param timeoutMs timeout value in milliseconds
* @param minScore the value in range (0.0, 1.0)
* @param resizeMode one of <em>ResizeMode</em> enum values.
* Set it to a value different from <em>NO_RESIZE</em>
* if the actual screenshot is expected to have different
* dimensions in comparison to the previously remembered one
* @return self instance for chaining
* @throws ScreenshotComparisonTimeout if the calculated score is still
* less than the given score after timeout happens
* @throws ScreenshotComparisonError if {@link #remember()} method has not been invoked yet
*/
public ScreenshotState verifyNotChanged(long timeoutMs, double minScore, ResizeMode resizeMode) {
return checkState((x) -> x >= minScore, timeoutMs, resizeMode);
}

private static Mat prepareImageForComparison(BufferedImage srcImage) {
final BufferedImage result = new BufferedImage(srcImage.getWidth(), srcImage.getHeight(),
BufferedImage.TYPE_3BYTE_BGR);
final Graphics2D g = srcImage.createGraphics();
try {
g.setComposite(AlphaComposite.Src);
g.drawImage(srcImage, 0, 0, null);
} finally {
g.dispose();
}
final byte[] pixels = ((DataBufferByte) result.getRaster().getDataBuffer()).getData();
final Mat imageMat = new Mat(result.getHeight(), result.getWidth(), CvType.CV_8UC3);
imageMat.put(0, 0, pixels);
return imageMat;
}

private static Mat resizeFirstMatrixToSecondMatrixResolution(Mat first, Mat second) {
if (first.width() != second.width() || first.height() != second.height()) {
final Mat result = new Mat();
final Size sz = new Size(second.width(), second.height());
Imgproc.resize(first, result, sz);
return result;
}
return first;
}

/**
* A shortcut to {@link #getOverlapScore(BufferedImage, BufferedImage, ResizeMode)} method
* for the case if both reference and template images are expected to have the same dimensions.
*
* @param refImage reference image
* @param tplImage template
* @return similarity score value in range (-1.0, 1.0). 1.0 is returned if the images are equal
* @throws ScreenshotComparisonError if provided images are not valid or have different resolution
*/
public static double getOverlapScore(BufferedImage refImage, BufferedImage tplImage) {
return getOverlapScore(refImage, tplImage, ResizeMode.NO_RESIZE);
}

/**
* Compares two valid java bitmaps and calculates similarity score between them.
*
* @param refImage reference image
* @param tplImage template
* @param resizeMode one of possible enum values. Set it either to <em>TEMPLATE_TO_REFERENCE_RESOLUTION</em> or
* <em>REFERENCE_TO_TEMPLATE_RESOLUTION</em> if given bitmaps have different dimensions
* @return similarity score value in range (-1.0, 1.0). 1.0 is returned if the images are equal
* @throws ScreenshotComparisonError if provided images are not valid or have
* different resolution, but resizeMode has been set to <em>NO_RESIZE</em>
*/
public static double getOverlapScore(BufferedImage refImage, BufferedImage tplImage, ResizeMode resizeMode) {
Mat ref = prepareImageForComparison(refImage);
if (ref.empty()) {
throw new ScreenshotComparisonError("Reference image cannot be converted for further comparison");
}
Mat tpl = prepareImageForComparison(tplImage);
if (tpl.empty()) {
throw new ScreenshotComparisonError("Template image cannot be converted for further comparison");
}
switch (resizeMode) {
case TEMPLATE_TO_REFERENCE_RESOLUTION:
tpl = resizeFirstMatrixToSecondMatrixResolution(tpl, ref);
break;
case REFERENCE_TO_TEMPLATE_RESOLUTION:
ref = resizeFirstMatrixToSecondMatrixResolution(ref, tpl);
break;
default:
// do nothing
}

if (ref.width() != tpl.width() || ref.height() != tpl.height()) {
throw new ScreenshotComparisonError(
"Resolutions of template and reference images are expected to be equal. "
+ "Try different resizeMode value."
);
}

final Mat res = new Mat(ref.rows() - tpl.rows() + 1, ref.cols() - tpl.cols() + 1, CvType.CV_32FC1);
Imgproc.matchTemplate(ref, tpl, res, Imgproc.TM_CCOEFF_NORMED);
Core.MinMaxLocResult minMaxLocResult = Core.minMaxLoc(res);
return minMaxLocResult.maxVal;
}

public enum ResizeMode {
NO_RESIZE, TEMPLATE_TO_REFERENCE_RESOLUTION, REFERENCE_TO_TEMPLATE_RESOLUTION
}
}
Loading

0 comments on commit d6a368c

Please sign in to comment.