diff --git a/build.gradle b/build.gradle
index 8ee52e4..ba27ad8 100644
--- a/build.gradle
+++ b/build.gradle
@@ -15,7 +15,7 @@ group 'com.github.roroche'
 if (project.hasProperty('newVersion')) {
     project.version = project.newVersion
 } else {
-    project.version = '0.0.1-SNAPSHOT'
+    project.version = '1.0.12-SNAPSHOT'
 }
 
 repositories {
@@ -30,6 +30,7 @@ dependencies {
     implementation localGroovy()
     implementation 'ch.ifocusit:plantuml-builder:1.4'
     implementation 'org.reflections:reflections:0.9.12'
+    implementation 'io.github.classgraph:classgraph:4.8.78'
     testImplementation('org.junit.jupiter:junit-jupiter:5.6.2')
     testImplementation 'com.pragmaticobjects.oo.tests:oo-tests:0.0.1'
     testImplementation 'com.pragmaticobjects.oo.tests:tests-junit5:0.0.1'
diff --git a/src/main/groovy/com/github/roroche/plantuml/PlantUmlPlugin.groovy b/src/main/groovy/com/github/roroche/plantuml/PlantUmlPlugin.groovy
index 43d40b8..5da0357 100644
--- a/src/main/groovy/com/github/roroche/plantuml/PlantUmlPlugin.groovy
+++ b/src/main/groovy/com/github/roroche/plantuml/PlantUmlPlugin.groovy
@@ -4,6 +4,8 @@ import com.github.roroche.plantuml.tasks.BuildClassDiagramTask
 import com.github.roroche.plantuml.tasks.ClassDiagramExtension
 import org.gradle.api.Plugin
 import org.gradle.api.Project
+import org.gradle.api.Task
+import org.gradle.api.UnknownTaskException
 
 /**
  * Custom Gradle plugin to generate and print PlantUML diagrams.
@@ -20,10 +22,28 @@ class PlantUmlPlugin implements Plugin<Project> {
                 "classDiagram",
                 ClassDiagramExtension.class
         )
-        project.tasks.create(
+        final BuildClassDiagramTask buildClassDiagramTask = project.tasks.create(
                 "buildClassDiagram",
                 BuildClassDiagramTask.class,
                 classDiagramExtension
         )
+        try {
+            final Task javaCompile = project.getTasks().getByName("compileJava")
+            buildClassDiagramTask.dependsOn(javaCompile)
+        } catch (final UnknownTaskException ignored) {
+            buildClassDiagramTask.getLogger().debug("Not a Java project")
+        }
+        try {
+            final Task groovyCompile = project.getTasks().getByName("compileGroovy")
+            buildClassDiagramTask.dependsOn(groovyCompile)
+        } catch (final UnknownTaskException ignored) {
+            buildClassDiagramTask.getLogger().debug("Not a Groovy project")
+        }
+        try {
+            final Task kotlinCompile = project.getTasks().getByName("compileKotlin")
+            buildClassDiagramTask.dependsOn(kotlinCompile)
+        } catch (final UnknownTaskException ignored) {
+            buildClassDiagramTask.getLogger().debug("Not a Kotlin project")
+        }
     }
 }
diff --git a/src/main/groovy/com/github/roroche/plantuml/tasks/BuildClassDiagramTask.groovy b/src/main/groovy/com/github/roroche/plantuml/tasks/BuildClassDiagramTask.groovy
index 9998d39..d326854 100644
--- a/src/main/groovy/com/github/roroche/plantuml/tasks/BuildClassDiagramTask.groovy
+++ b/src/main/groovy/com/github/roroche/plantuml/tasks/BuildClassDiagramTask.groovy
@@ -6,6 +6,7 @@ import com.github.roroche.plantuml.diagrams.Diagram
 import com.github.roroche.plantuml.diagrams.DiagramWithLog
 import org.gradle.api.DefaultTask
 import org.gradle.api.tasks.TaskAction
+import org.reflections.Reflections
 
 import javax.inject.Inject
 
@@ -30,17 +31,35 @@ class BuildClassDiagramTask extends DefaultTask implements CustomTask {
     @Override
     void execute() {
         getLogger().debug(
-                String.format("Package to scan: %s", extension.packageName)
+                "Package to scan: {}",
+                extension.packageName
         )
         getLogger().debug(
-                String.format("Output file: %s", extension.outputFile)
+                "Output file: {}",
+                extension.outputFile
         )
         getLogger().debug(
-                String.format("Classes to ignore: %s", extension.ignoredClasses)
+                "Classes to ignore: {}",
+                extension.ignoredClasses
         )
+        final URL[] urls = project.sourceSets.main.output.classesDirs.files.collect { File dir ->
+            dir.listFiles()
+        }.flatten().collect {
+            it.toURI().toURL()
+        } as URL[]
+        getLogger().lifecycle(
+                "URLs to scan: " + urls
+        )
+        final ClassLoader classLoader = new URLClassLoader(urls)
         final Classes classes = new ClsWithLog(
                 new ClsFiltered(
-                        new ClsInPackage(extension.packageName),
+                        new ClsInPackage(
+                                extension.packageName,
+                                new Reflections(
+                                        extension.packageName,
+                                        classLoader
+                                )
+                        ),
                         new ClsWithNames(extension.ignoredClasses)
                 ),
                 getLogger()
diff --git a/src/main/kotlin/com/github/roroche/plantuml/classes/ClsInPackage.kt b/src/main/kotlin/com/github/roroche/plantuml/classes/ClsInPackage.kt
index 85df4ed..dfc25ea 100644
--- a/src/main/kotlin/com/github/roroche/plantuml/classes/ClsInPackage.kt
+++ b/src/main/kotlin/com/github/roroche/plantuml/classes/ClsInPackage.kt
@@ -1,50 +1,32 @@
 package com.github.roroche.plantuml.classes
 
 import com.github.roroche.plantuml.classes.exceptions.InvalidPackageException
+import org.reflections.Configuration
 import org.reflections.Reflections
-import org.reflections.scanners.SubTypesScanner
-import org.reflections.scanners.TypeAnnotationsScanner
-import org.reflections.util.ClasspathHelper
-import org.reflections.util.ConfigurationBuilder
-import java.net.URL
-import java.util.concurrent.Executors
 
 /**
  * Utility class to find [Classes] in a given package.
  *
  * @property packageName The name of the package to scan.
- * @property packageUrls URL of the package (passed as [Collection]).
  * @property reflections Utility to find classes in pckage.
  */
 class ClsInPackage(
     private val packageName: String,
-    private val packageUrls: Collection<URL>,
     private val reflections: Reflections
 ) : Classes {
 
     /**
-     * Secondary constructor using URLs.
+     * Secondary constructor with reflections configuration.
      *
      * @param packageName The name of the package to scan.
-     * @param packageUrls URL of the package (passed as [Collection]).
+     * @param configuration The [Configuration] to use.
      */
     constructor(
         packageName: String,
-        packageUrls: Collection<URL>
+        configuration: Configuration
     ) : this(
-        packageName,
-        packageUrls,
-        Reflections(
-            ConfigurationBuilder()
-                .setUrls(
-                    packageUrls
-                ).setScanners(
-                    SubTypesScanner(false),
-                    TypeAnnotationsScanner()
-                ).setExecutorService(
-                    Executors.newFixedThreadPool(4)
-                )
-        )
+        packageName = packageName,
+        reflections = Reflections(configuration)
     )
 
     /**
@@ -52,20 +34,26 @@ class ClsInPackage(
      *
      * @param packageName The name of the package to scan.
      */
-    constructor(packageName: String) : this(
-        packageName,
-        ClasspathHelper.forPackage(packageName)
+    constructor(
+        packageName: String
+    ) : this(
+        packageName = packageName,
+        reflections = Reflections(
+            packageName,
+            Thread.currentThread().contextClassLoader
+        )
     )
 
     /**
      * @return Classes to be used for diagram generation.
      */
     override fun list(): List<Class<out Any>> {
-        if (packageUrls.isNullOrEmpty()) {
-            throw InvalidPackageException(packageName)
-        }
-        return reflections.getSubTypesOf(
+        val list = reflections.getSubTypesOf(
             Any::class.java
         ).asIterable().toList()
+        if (list.isNullOrEmpty()) {
+            throw InvalidPackageException(packageName)
+        }
+        return list
     }
 }
\ No newline at end of file