Merge pull request #2 from keshihoriuchi/using-multi-core

Using thread pool to use multi core
keshihoriuchi · May 23, 2021 · 5a6f43f · 5a6f43f
2 parents 993bcb1 + 3b56474
commit 5a6f43f
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 17 deletions.
diff --git a/cli/src/main/java/io/github/keshihoriuchi/gahi/Cli.java b/cli/src/main/java/io/github/keshihoriuchi/gahi/Cli.java
@@ -17,7 +17,7 @@
 public class Cli {
     private static ImageSD imageSD;
 
-    public static void main(String[] args) throws IOException {
+    public static void main(String[] args) throws IOException, InterruptedException {
         CommandLineParser parser = new DefaultParser();
         final Options options = new Options();
         options.addOption(new Option("a", "algorithm", true, "Algorithm to classify image"));
@@ -58,7 +58,7 @@ else if (parsedArgs[0].equals("search")) {
         }
     }
 
-    static void prepare(String dir, Class algo) throws IOException {
+    static void prepare(String dir, Class algo) throws IOException, InterruptedException {
         imageSD = new ImageSD(Paths.get(dir), Paths.get(".", "./temp_index"), algo);
         imageSD.clearIndex();
         imageSD.createIndex();

diff --git a/cli/src/main/java/io/github/keshihoriuchi/gahi/ImageSD.java b/cli/src/main/java/io/github/keshihoriuchi/gahi/ImageSD.java
@@ -7,6 +7,9 @@
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import javax.imageio.ImageIO;
@@ -93,28 +96,21 @@ public Stream<? extends List<String>> dup() throws IOException {
         return result;
     }
 
-    public void createIndex() throws IOException {
+    public void createIndex() throws IOException, InterruptedException {
         List<String> images = FileUtils.getAllImages(baseDir, isRecursive);
         int size = images.size();
         DocumentBuilder builder = new GlobalDocumentBuilder(algo);
 
+        ExecutorService pool = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
         try (IndexWriter iw = LuceneUtils.createIndexWriter(FSDirectory.open(indexDir), true, LuceneUtils.AnalyzerType.WhitespaceAnalyzer)) {
+            List<Callable<Void>> tasks = new ArrayList<>(size);
             for (int i = 0; i < size; i++) {
-                String imgPath = images.get(i);
-                try {
-                    BufferedImage img = ImageIO.read(new File(imgPath));
-                    Document document = builder.createDocument(img, imgPath);
-                    iw.addDocument(document);
-                    JSONObject jo = new JSONObject();
-                    jo.put("type", "index_creating");
-                    jo.put("path", imgPath);
-                    jo.put("total", size);
-                    jo.put("finished", i + 1);
-                    System.out.println(jo);
-                } catch (RuntimeException e) {
-                    System.err.println(e);
-                }
+                Callable task = new CreateIndexTask(images.get(i), builder, iw, size);
+                tasks.add(task);
             }
+            pool.invokeAll(tasks);
+        } finally {
+            pool.shutdown();
         }
     }
 
@@ -131,4 +127,38 @@ public void printDocs() throws IOException {
             }
         }
     }
+
+    static class CreateIndexTask implements Callable<Void> {
+        final String imgPath;
+        final DocumentBuilder builder;
+        final IndexWriter iw;
+        final int size;
+        public CreateIndexTask(String imgPath, DocumentBuilder builder, IndexWriter iw, int size) {
+            this.imgPath = imgPath;
+            this.builder = builder;
+            this.iw = iw;
+            this.size = size;
+        }
+
+        public Void call() {
+            try {
+                BufferedImage img = ImageIO.read(new File(imgPath));
+                Document document = builder.createDocument(img, imgPath);
+                iw.addDocument(document);
+                JSONObject jo = new JSONObject();
+                jo.put("type", "index_creating");
+                jo.put("path", imgPath);
+                jo.put("total", size);
+                safePrintln(jo);
+            } catch (RuntimeException | IOException e) {
+                System.err.println(e);
+            }
+            return null;
+        }
+        private void safePrintln(JSONObject s) {
+            synchronized (System.out) {
+                System.out.println(s);
+            }
+        }
+    }
 }
diff --git a/gui/src/main.ts b/gui/src/main.ts
@@ -77,6 +77,7 @@ ipcMain.on("start-cli", async (ev, dirPath, algo) => {
   console.log(cmd);
   cli = child_process.spawn("cmd.exe", ["/C", cmd, "dup", "-a", algo, dirPath]);
   let errStr = "";
+  let finished = 0;
   cli.stdout
     .pipe(iconv.decodeStream(`cp${cpCode}`))
     .pipe(split2())
@@ -90,6 +91,8 @@ ipcMain.on("start-cli", async (ev, dirPath, algo) => {
         return;
       }
       if (d.type === "index_creating") {
+        finished++;
+        d.finished = finished;
         ev.reply("cli-interm", d);
         mainWindow.setProgressBar(d.finished / d.total);
       } else if (d.type === "finish") {