1/*
2 * Copyright (c) 1997, 2016, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation.
8 *
9 * This code is distributed in the hope that it will be useful, but WITHOUT
10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
12 * version 2 for more details (a copy is included in the LICENSE file that
13 * accompanied this code).
14 *
15 * You should have received a copy of the GNU General Public License version
16 * 2 along with this work; if not, write to the Free Software Foundation,
17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18 *
19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20 * or visit www.oracle.com if you need additional information or have any
21 * questions.
22 */
23package org.netbeans.jemmy.util;
24
25import java.awt.AWTException;
26import java.awt.Component;
27import java.awt.Rectangle;
28import java.awt.Robot;
29import java.awt.Toolkit;
30import java.awt.image.BufferedImage;
31import java.io.BufferedOutputStream;
32import java.io.ByteArrayOutputStream;
33import java.io.FileOutputStream;
34import java.io.IOException;
35import java.io.OutputStream;
36import java.util.zip.CRC32;
37import java.util.zip.Deflater;
38import java.util.zip.DeflaterOutputStream;
39
40/**
41 * This class allows to encode BufferedImage into B/W, greyscale or true color
42 * PNG image format with maximum compression.<br>
43 * It also provides complete functionality for capturing full screen, part of
44 * screen or single component, encoding and saving captured image info PNG file.
45 *
46 * @author Adam Sotona
47 * @version 1.0
48 */
49public class PNGEncoder extends Object {
50
51    /**
52     * black and white image mode.
53     */
54    public static final byte BW_MODE = 0;
55    /**
56     * grey scale image mode.
57     */
58    public static final byte GREYSCALE_MODE = 1;
59    /**
60     * full color image mode.
61     */
62    public static final byte COLOR_MODE = 2;
63
64    OutputStream out;
65    CRC32 crc;
66    byte mode;
67
68    /**
69     * public constructor of PNGEncoder class with greyscale mode by default.
70     *
71     * @param out output stream for PNG image format to write into
72     */
73    public PNGEncoder(OutputStream out) {
74        this(out, GREYSCALE_MODE);
75    }
76
77    /**
78     * public constructor of PNGEncoder class.
79     *
80     * @param out output stream for PNG image format to write into
81     * @param mode BW_MODE, GREYSCALE_MODE or COLOR_MODE
82     */
83    public PNGEncoder(OutputStream out, byte mode) {
84        crc = new CRC32();
85        this.out = out;
86        if (mode < 0 || mode > 2) {
87            throw new IllegalArgumentException("Unknown color mode");
88        }
89        this.mode = mode;
90    }
91
92    void write(int i) throws IOException {
93        byte b[] = {(byte) ((i >> 24) & 0xff), (byte) ((i >> 16) & 0xff), (byte) ((i >> 8) & 0xff), (byte) (i & 0xff)};
94        write(b);
95    }
96
97    void write(byte b[]) throws IOException {
98        out.write(b);
99        crc.update(b);
100    }
101
102    /**
103     * main encoding method (stays blocked till encoding is finished).
104     *
105     * @param image BufferedImage to encode
106     * @throws IOException IOException
107     */
108    public void encode(BufferedImage image) throws IOException {
109        int width = image.getWidth(null);
110        int height = image.getHeight(null);
111        final byte id[] = {-119, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13};
112        write(id);
113        crc.reset();
114        write("IHDR".getBytes());
115        write(width);
116        write(height);
117        byte head[];
118        switch (mode) {
119            case BW_MODE:
120                head = new byte[]{1, 0, 0, 0, 0};
121                break;
122            case GREYSCALE_MODE:
123                head = new byte[]{8, 0, 0, 0, 0};
124                break;
125            case COLOR_MODE:
126                head = new byte[]{8, 2, 0, 0, 0};
127                break;
128            default:
129                throw new AssertionError("Unexpected mode: " + mode);
130        }
131        write(head);
132        write((int) crc.getValue());
133        ByteArrayOutputStream compressed = new ByteArrayOutputStream(65536);
134        BufferedOutputStream bos = new BufferedOutputStream(new DeflaterOutputStream(compressed, new Deflater(9)));
135        int pixel;
136        int color;
137        int colorset;
138        switch (mode) {
139            case BW_MODE:
140                int rest = width % 8;
141                int bytes = width / 8;
142                for (int y = 0; y < height; y++) {
143                    bos.write(0);
144                    for (int x = 0; x < bytes; x++) {
145                        colorset = 0;
146                        for (int sh = 0; sh < 8; sh++) {
147                            pixel = image.getRGB(x * 8 + sh, y);
148                            color = ((pixel >> 16) & 0xff);
149                            color += ((pixel >> 8) & 0xff);
150                            color += (pixel & 0xff);
151                            colorset <<= 1;
152                            if (color >= 3 * 128) {
153                                colorset |= 1;
154                            }
155                        }
156                        bos.write((byte) colorset);
157                    }
158                    if (rest > 0) {
159                        colorset = 0;
160                        for (int sh = 0; sh < width % 8; sh++) {
161                            pixel = image.getRGB(bytes * 8 + sh, y);
162                            color = ((pixel >> 16) & 0xff);
163                            color += ((pixel >> 8) & 0xff);
164                            color += (pixel & 0xff);
165                            colorset <<= 1;
166                            if (color >= 3 * 128) {
167                                colorset |= 1;
168                            }
169                        }
170                        colorset <<= 8 - rest;
171                        bos.write((byte) colorset);
172                    }
173                }
174                break;
175            case GREYSCALE_MODE:
176                for (int y = 0; y < height; y++) {
177                    bos.write(0);
178                    for (int x = 0; x < width; x++) {
179                        pixel = image.getRGB(x, y);
180                        color = ((pixel >> 16) & 0xff);
181                        color += ((pixel >> 8) & 0xff);
182                        color += (pixel & 0xff);
183                        bos.write((byte) (color / 3));
184                    }
185                }
186                break;
187            case COLOR_MODE:
188                for (int y = 0; y < height; y++) {
189                    bos.write(0);
190                    for (int x = 0; x < width; x++) {
191                        pixel = image.getRGB(x, y);
192                        bos.write((byte) ((pixel >> 16) & 0xff));
193                        bos.write((byte) ((pixel >> 8) & 0xff));
194                        bos.write((byte) (pixel & 0xff));
195                    }
196                }
197                break;
198        }
199        bos.close();
200        write(compressed.size());
201        crc.reset();
202        write("IDAT".getBytes());
203        write(compressed.toByteArray());
204        write((int) crc.getValue());
205        write(0);
206        crc.reset();
207        write("IEND".getBytes());
208        write((int) crc.getValue());
209        out.close();
210    }
211
212    /**
213     * Static method performing screen capture into PNG image format file with
214     * given fileName.
215     *
216     * @param rect Rectangle of screen to be captured
217     * @param fileName file name for screen capture PNG image file
218     */
219    public static void captureScreen(Rectangle rect, String fileName) {
220        captureScreen(rect, fileName, GREYSCALE_MODE);
221    }
222
223    /**
224     * Static method performing screen capture into PNG image format file with
225     * given fileName.
226     *
227     * @param rect Rectangle of screen to be captured
228     * @param mode image color mode
229     * @param fileName file name for screen capture PNG image file
230     */
231    public static void captureScreen(Rectangle rect, String fileName, byte mode) {
232        try {
233            BufferedImage capture = new Robot().createScreenCapture(rect);
234            BufferedOutputStream file = new BufferedOutputStream(new FileOutputStream(fileName));
235            PNGEncoder encoder = new PNGEncoder(file, mode);
236            encoder.encode(capture);
237        } catch (AWTException awte) {
238            awte.printStackTrace();
239        } catch (IOException ioe) {
240            ioe.printStackTrace();
241        }
242    }
243
244    /**
245     * Static method performing one component screen capture into PNG image
246     * format file with given fileName.
247     *
248     * @param comp Component to be captured
249     * @param fileName String image target filename
250     */
251    public static void captureScreen(Component comp, String fileName) {
252        captureScreen(comp, fileName, GREYSCALE_MODE);
253    }
254
255    /**
256     * Static method performing one component screen capture into PNG image
257     * format file with given fileName.
258     *
259     * @param comp Component to be captured
260     * @param fileName String image target filename
261     * @param mode image color mode
262     */
263    public static void captureScreen(Component comp, String fileName, byte mode) {
264        captureScreen(new Rectangle(comp.getLocationOnScreen(),
265                comp.getSize()),
266                fileName, mode);
267    }
268
269    /**
270     * Static method performing whole screen capture into PNG image format file
271     * with given fileName.
272     *
273     * @param fileName String image target filename
274     */
275    public static void captureScreen(String fileName) {
276        captureScreen(fileName, GREYSCALE_MODE);
277    }
278
279    /**
280     * Static method performing whole screen capture into PNG image format file
281     * with given fileName.
282     *
283     * @param fileName String image target filename
284     * @param mode image color mode
285     */
286    public static void captureScreen(String fileName, byte mode) {
287        captureScreen(new Rectangle(Toolkit.getDefaultToolkit().getScreenSize()), fileName, mode);
288    }
289}
290