diff --git a/src/Analysis.php b/src/Analysis.php
index a0fe8384..f152f838 100644
--- a/src/Analysis.php
+++ b/src/Analysis.php
@@ -7,7 +7,6 @@
 namespace OpenApi;
 
 use OpenApi\Annotations as OA;
-use OpenApi\Attributes as OAT;
 use OpenApi\Processors\ProcessorInterface;
 
 /**
@@ -339,7 +338,7 @@ public function getSchemaForSource(string $fqdn): ?OA\Schema
                 $definition = $definitions[$fqdn];
                 if (is_iterable($definition['context']->annotations)) {
                     foreach (array_reverse($definition['context']->annotations) as $annotation) {
-                        if (in_array(get_class($annotation), [OA\Schema::class, OAT\Schema::class]) && !$annotation->_context->is('generated')) {
+                        if ($annotation->isRoot(OA\Schema::class) && !$annotation->_context->is('generated')) {
                             return $annotation;
                         }
                     }