diff --git a/puml_generator.py b/puml_generator.py
index 93be27c..8f06e01 100644
--- a/puml_generator.py
+++ b/puml_generator.py
@@ -204,7 +204,7 @@ def get_method_name_from_line(line:str) -> str:
     return line.strip().split("(")[0].split(" ")[-1]
 
 
-def class_to_puml(filename:str) -> JavaClass:
+def class_to_puml(filename:str, base_package_slug:str) -> JavaClass:
     imported_classes = {}
     puml_content = []
     uml_relations = {}
@@ -243,7 +243,7 @@ def class_to_puml(filename:str) -> JavaClass:
                     class_methods[method_name] = get_method_visibility(javaline)
                 next_line_setter = False
             elif javaline.startswith("import"):
-                if "*" not in javaline:
+                if "*" not in javaline and base_package_slug in javaline:
                     importline_package_slug = javaline.strip().split(" ")[1].replace(";", "")
                     imported_classes[importline_package_slug] = importline_package_slug
 
@@ -319,7 +319,7 @@ def generate_puml(root_dir:str=".", no_pkgs:bool=False) -> None:
     puml_code = ""
 
     for class_file in java_files:
-        puml_code += class_to_puml(class_file).get_puml_content()
+        puml_code += class_to_puml(class_file, base_package_slug).get_puml_content()
 
     # write the PlantUML code to a file
     with open(f"{root_dir}/generated_class_diagram.puml", "w") as puml_file: