Skip to content

Commit

Permalink
Added multithreaded feature for image extraction (#1641)
Browse files Browse the repository at this point in the history 
Co-authored-by: Anthony Stirling <[email protected]>
  • Loading branch information
@HimaGirija99 @Frooodle
HimaGirija99 and Frooodle authored Aug 7, 2024
1 parent 96b0805 commit e50391a
Showing 1 changed file with 121 additions and 73 deletions.
194 changes: 121 additions & 73 deletions src/main/java/stirling/software/SPDF/controller/api/misc/ExtractImagesController.java
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,8 @@
package stirling.software.SPDF.controller.api.misc;

import java.awt.Graphics2D;
import java.awt.Image;
import java.awt.image.BufferedImage;
import java.awt.image.RenderedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.zip.Deflater;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import javax.imageio.ImageIO;

import io.github.pixee.security.Filenames;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.pdmodel.PDDocument;
Expand All @@ -28,14 +17,25 @@
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import io.github.pixee.security.Filenames;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;

import stirling.software.SPDF.model.api.PDFWithImageFormatRequest;
import stirling.software.SPDF.utils.WebResponseUtils;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.RenderedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.zip.Deflater;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

@RestController
@RequestMapping("/api/v1/misc")
@Tag(name = "Misc", description = "Miscellaneous APIs")
Expand All @@ -47,16 +47,19 @@ public class ExtractImagesController {
@Operation(
summary = "Extract images from a PDF file",
description =
"This endpoint extracts images from a given PDF file and returns them in a zip file. Users can specify the output image format. Input:PDF Output:IMAGE/ZIP Type:SIMO")
"This endpoint extracts images from a given PDF file and returns them in a zip file. Users can specify the output image format. Input: PDF Output: IMAGE/ZIP Type: SIMO")
public ResponseEntity<byte[]> extractImages(@ModelAttribute PDFWithImageFormatRequest request)
throws IOException {
throws IOException, InterruptedException, ExecutionException {
MultipartFile file = request.getFileInput();
String format = request.getFormat();

System.out.println(
System.currentTimeMillis() + "file=" + file.getName() + ", format=" + format);
System.currentTimeMillis() + " file=" + file.getName() + ", format=" + format);
PDDocument document = Loader.loadPDF(file.getBytes());

// Determine if multithreading should be used based on PDF size or number of pages
boolean useMultithreading = shouldUseMultithreading(file, document);

// Create ByteArrayOutputStream to write zip file to byte array
ByteArrayOutputStream baos = new ByteArrayOutputStream();

Expand All @@ -66,76 +69,121 @@ public ResponseEntity<byte[]> extractImages(@ModelAttribute PDFWithImageFormatRe
// Set compression level
zos.setLevel(Deflater.BEST_COMPRESSION);

int imageIndex = 1;
String filename =
Filenames.toSimpleFileName(file.getOriginalFilename())
.replaceFirst("[.][^.]+$", "");
int pageNum = 0;
Set<Integer> processedImages = new HashSet<>();
// Iterate over each page
for (PDPage page : document.getPages()) {
++pageNum;
// Extract images from page
for (COSName name : page.getResources().getXObjectNames()) {
if (page.getResources().isImageXObject(name)) {
PDImageXObject image = (PDImageXObject) page.getResources().getXObject(name);
int imageHash = image.hashCode();

if (useMultithreading) {
// Executor service to handle multithreading
ExecutorService executor =
Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
Set<Future<Void>> futures = new HashSet<>();

// Iterate over each page
for (int pgNum = 0; pgNum < document.getPages().getCount(); pgNum++) {
PDPage page = document.getPage(pgNum);
int pageNum = document.getPages().indexOf(page) + 1;
// Submit a task for processing each page
Future<Void> future =
executor.submit(
() -> {
extractImagesFromPage(
page, format, filename, pageNum, processedImages, zos);
return null;
});

futures.add(future);
}

// Wait for all tasks to complete
for (Future<Void> future : futures) {
future.get();
}

// Close executor service
executor.shutdown();
} else {
// Single-threaded extraction
for (int pgNum = 0; pgNum < document.getPages().getCount(); pgNum++) {
PDPage page = document.getPage(pgNum);
extractImagesFromPage(page, format, filename, pgNum + 1, processedImages, zos);
}
}

// Close PDDocument and ZipOutputStream
document.close();
zos.close();

// Create ByteArrayResource from byte array
byte[] zipContents = baos.toByteArray();

return WebResponseUtils.boasToWebResponse(
baos, filename + "_extracted-images.zip", MediaType.APPLICATION_OCTET_STREAM);
}

private boolean shouldUseMultithreading(MultipartFile file, PDDocument document) {
// Criteria: Use multithreading if file size > 10MB or number of pages > 20
long fileSizeInMB = file.getSize() / (1024 * 1024);
int numberOfPages = document.getPages().getCount();
return fileSizeInMB > 10 || numberOfPages > 20;
}

private void extractImagesFromPage(
PDPage page,
String format,
String filename,
int pageNum,
Set<Integer> processedImages,
ZipOutputStream zos)
throws IOException {
for (COSName name : page.getResources().getXObjectNames()) {
if (page.getResources().isImageXObject(name)) {
PDImageXObject image = (PDImageXObject) page.getResources().getXObject(name);
int imageHash = image.hashCode();
synchronized (processedImages) {
if (processedImages.contains(imageHash)) {
continue; // Skip already processed images
}
processedImages.add(imageHash);
}

// Convert image to desired format
RenderedImage renderedImage = image.getImage();
BufferedImage bufferedImage = null;
if ("png".equalsIgnoreCase(format)) {
bufferedImage =
new BufferedImage(
renderedImage.getWidth(),
renderedImage.getHeight(),
BufferedImage.TYPE_INT_ARGB);
} else if ("jpeg".equalsIgnoreCase(format) || "jpg".equalsIgnoreCase(format)) {
bufferedImage =
new BufferedImage(
renderedImage.getWidth(),
renderedImage.getHeight(),
BufferedImage.TYPE_INT_RGB);
} else if ("gif".equalsIgnoreCase(format)) {
bufferedImage =
new BufferedImage(
renderedImage.getWidth(),
renderedImage.getHeight(),
BufferedImage.TYPE_BYTE_INDEXED);
}
RenderedImage renderedImage = image.getImage();

// Write image to zip file
String imageName =
filename + "_" + imageIndex + " (Page " + pageNum + ")." + format;
ZipEntry zipEntry = new ZipEntry(imageName);
zos.putNextEntry(zipEntry);
// Convert to standard RGB colorspace if needed
BufferedImage bufferedImage = convertToRGB(renderedImage, format);

Graphics2D g = bufferedImage.createGraphics();
g.drawImage((Image) renderedImage, 0, 0, null);
g.dispose();
// Write image bytes to zip file
// Write image to zip file
String imageName = filename + "_" + imageHash + " (Page " + pageNum + ")." + format;
synchronized (zos) {
zos.putNextEntry(new ZipEntry(imageName));
ByteArrayOutputStream imageBaos = new ByteArrayOutputStream();
ImageIO.write(bufferedImage, format, imageBaos);
zos.write(imageBaos.toByteArray());

zos.closeEntry();
imageIndex++;
}
}
}
}

// Close ZipOutputStream and PDDocument
zos.close();
document.close();

// Create ByteArrayResource from byte array
byte[] zipContents = baos.toByteArray();
private BufferedImage convertToRGB(RenderedImage renderedImage, String format) {
int width = renderedImage.getWidth();
int height = renderedImage.getHeight();
BufferedImage rgbImage;

if ("png".equalsIgnoreCase(format)) {
rgbImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
} else if ("jpeg".equalsIgnoreCase(format) || "jpg".equalsIgnoreCase(format)) {
rgbImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
} else if ("gif".equalsIgnoreCase(format)) {
rgbImage = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_INDEXED);
} else {
rgbImage = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
}

return WebResponseUtils.boasToWebResponse(
baos, filename + "_extracted-images.zip", MediaType.APPLICATION_OCTET_STREAM);
Graphics2D g = rgbImage.createGraphics();
g.drawImage((Image) renderedImage, 0, 0, null);
g.dispose();
return rgbImage;
}
}

0 comments on commit e50391a

Please sign in to comment.