source: gs3-extensions/atea-nlp-tools/trunk/src/ocr/src/main/java/org/atea/nlptools/ocr/servlets/TesseractServlet.java@ 35733

package org.atea.nlptools.ocr.servlets;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Properties;
import java.util.UUID;

import javax.servlet.ServletException;
import javax.servlet.annotation.MultipartConfig;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.Part;

import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import com.google.gson.stream.JsonWriter;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import org.atea.nlptools.ocr.abstractions.objects.HttpStatusCode;
import org.atea.nlptools.ocr.abstractions.objects.IOcrFile;
import org.atea.nlptools.ocr.abstractions.objects.IOcrOptions;
import org.atea.nlptools.ocr.abstractions.services.IOcrService;
import org.atea.nlptools.ocr.objects.OcrFile;
import org.atea.nlptools.ocr.objects.OcrOptionsDto;
import org.atea.nlptools.ocr.objects.TesseractOptions;
import org.atea.nlptools.ocr.services.TesseractOcrService;
import org.bytedeco.tesseract.global.tesseract;

@MultipartConfig
public class TesseractServlet extends HttpServlet
{
    /**
     * Gets the maximum number of parts that may be submitted in any one request.
     */
    private static final int MaxParts = 10;

    /**
     * Gets the maximum size in bytes that a part may be.
     */
    private static final int MaxPartSize = 1024 * 1024 * 50;

    private static final Logger logger = LogManager.getLogger(TesseractServlet.class);
    private static final Type optionsTypeToken = new TypeToken<HashMap<String, OcrOptionsDto>>() {}.getType();

    private final Gson gsonInstance;

    private IOcrService ocrService;

    public TesseractServlet()
    {
        gsonInstance = new Gson();
    }

    @Override
    public void init() 
    {
        Properties prop = new Properties();

        try
        {
            FileInputStream fis = new FileInputStream("../webapps/gs3-atea-ocr/config.properties");
            prop.load(fis);
            fis.close();
        }
        catch (IOException ex) 
        {
            logger.error("Failed to retrieve properties file", ex);
        }

        String tessBin = prop.getProperty("tesseract.data.path");
        this.ocrService = new TesseractOcrService(tessBin);

        logger.info("Initialised!");
    }

    @Override
    protected void doGet(HttpServletRequest request, HttpServletResponse response)
        throws IOException
    {
        response.sendError(HttpStatusCode.ClientForbidden, "POST Multipart request expected.");
    }
    
    @Override
    protected void doPost(HttpServletRequest request, HttpServletResponse response)
            throws ServletException, IOException
    {
        try
        {
            logger.trace("POST request received.");
            final UUID requestUUID = UUID.randomUUID();

            final TesseractOptions genericOptions = new TesseractOptions();
            HashMap<String, IOcrOptions> options = ParseOptions(request);
            
            final List<IOcrFile> files = CreatePartTempFiles(request.getParts(), requestUUID, response);
            if (files == null) {
                return;
            }

            response.setContentType("application/json; charset=UTF-8");
            JsonWriter writer = gsonInstance.newJsonWriter(response.getWriter());
            writer.beginArray();

            for (IOcrFile file : files)
            {
                IOcrOptions actualOptions = options.get(file.getKey());
                if (actualOptions == null) {
                    actualOptions = genericOptions;
                }

                String ocrOutput = ocrService.run(file.getTempFile(), actualOptions);

                writer.beginObject();

                writer.name("key");
                writer.value(file.getKey());

                writer.name("fileName");
                writer.value(file.getFileName());

                writer.name("text");
                writer.value(ocrOutput);

                writer.endObject();
            }

            writer.endArray();
            writer.flush();
        }
        catch (Exception ex)
        {
            response.sendError(HttpStatusCode.ServerInternalServerError, "Failed to process the request.");
            logger.error("Failed to complete API call", ex);
        }
    }
    
    /**
     * Saves the parts in a request to temporary files. 
     * @param parts The parts.
     * @param requestUUID The assigned UUID of the request.
     * @param response The response.
     * @return The list of temporary files, or <code>null</code> if the operation failed.
     * @throws IOException
     * @throws ServletException
     */
    private List<IOcrFile> CreatePartTempFiles(Collection<Part> parts, UUID requestUUID, HttpServletResponse response)
        throws IOException, ServletException
    {
        if (parts.size() > MaxParts)
        {
            response.sendError(HttpStatusCode.ClientPayloadTooLarge, "No more than " + MaxParts + " parts may be submitted.");
            return null;
        }

        ArrayList<IOcrFile> files = new ArrayList<>();

        for (Part p : parts)
        {
            if (p.getSubmittedFileName() == null) {
                continue;
            }

            if (p.getSize() > MaxPartSize)
            {
                response.sendError(HttpStatusCode.ClientPayloadTooLarge, "A submitted part must be no more than " + MaxPartSize + " bytes");
                return null;
            }

            File tempFile = File.createTempFile("atea-ocr", requestUUID.toString());
            p.write(tempFile.getAbsolutePath());

            files.add(new OcrFile(p.getName(), p.getSubmittedFileName(), tempFile));
        }
        
        return files;
    }

    /**
     * Parses the options parameter.
     * @param request The request.
     * @return The submitted options.
     */
    private HashMap<String, IOcrOptions> ParseOptions(HttpServletRequest request)
    {
        HashMap<String, IOcrOptions> ret = new HashMap<>();

        String options = request.getParameter("options");
        if (options == null) {
            return ret;
        }

        HashMap<String, OcrOptionsDto> dtoOptions = gsonInstance.fromJson(options, optionsTypeToken);

        for (String key : dtoOptions.keySet())
        {
            OcrOptionsDto dto = dtoOptions.get(key);
            TesseractOptions tOptions = new TesseractOptions();

            if (dto.layoutDetection) {
                tOptions.pageSegmentationMode = tesseract.PSM_AUTO_OSD;
            }

            ret.put(key, tOptions);
        }
        
        return ret;
    }
}