diff --git a/google_vision.module b/google_vision.module
index 580ac94..dea6244 100644
--- a/google_vision.module
+++ b/google_vision.module
@@ -10,6 +10,8 @@ use Drupal\google_vision\GoogleVisionAPI;
 use Drupal\field\FieldStorageConfigInterface;
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Entity\EntityTypeInterface;
+use Drupal\Core\Field\FieldItemListInterface;
+use Drupal\Core\Field\FieldDefinitionInterface;
 
 /**
  * Implements hook_entity_presave().
@@ -76,11 +78,16 @@ function google_vision_form_field_config_form_taxonomy_builder($entity_type, Fie
 }
 
 /**
- * Implements hook_entity_type_alter().
+ * Implements hook_entity_bundle_field_info_alter().
  */
-function google_vision_entity_type_alter(array &$entity_types) {
-  $node = $entity_types['node'];
-  $node->addConstraint('SafeSearch');
+function google_vision_entity_bundle_field_info_alter(&$fields, EntityTypeInterface $entity_type, $bundle) {
+  if ($entity_type->id() == 'node') {
+    foreach ($fields as $field) {
+      if ($field->getType() == 'image') {
+        $fields[$field->getName()]->addConstraint('SafeSearch');
+      }
+    }
+  }
 }
 
 /**
diff --git a/src/Plugin/Validation/Constraint/SafeSearchConstraint.php b/src/Plugin/Validation/Constraint/SafeSearchConstraint.php
index aa40208..9bf064d 100755
--- a/src/Plugin/Validation/Constraint/SafeSearchConstraint.php
+++ b/src/Plugin/Validation/Constraint/SafeSearchConstraint.php
@@ -14,5 +14,5 @@ use Symfony\Component\Validator\Constraint;
  */
 class SafeSearchConstraint extends Constraint {
 
-  public $message = 'The file %title contains explicit content and will not be saved.';
+  public $message = 'This image contains explicit content and will not be saved.';
 }
diff --git a/src/Plugin/Validation/Constraint/SafeSearchConstraintValidator.php b/src/Plugin/Validation/Constraint/SafeSearchConstraintValidator.php
index 14881fc..02e59d3 100755
--- a/src/Plugin/Validation/Constraint/SafeSearchConstraintValidator.php
+++ b/src/Plugin/Validation/Constraint/SafeSearchConstraintValidator.php
@@ -4,16 +4,13 @@ namespace Drupal\google_vision\Plugin\Validation\Constraint;
 
 use Symfony\Component\Validator\Constraint;
 use Symfony\Component\Validator\ConstraintValidator;
-use Drupal\google_vision\GoogleVisionAPI;
 use Drupal\Core\Field\FieldItemListInterface;
-use Drupal\file\FileInterface;
-use Drupal\node\NodeInterface;
-use Drupal\field\FieldConfigInterface;
-use Drupal\field\FieldStorageConfigInterface;
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Entity\EntityTypeInterface;
+use Drupal\file\FileInterface;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 use Drupal\Core\DependencyInjection\ContainerInjectionInterface;
+use Drupal\google_vision\GoogleVisionAPI;
 
 /**
  * Validates the SafeSearch constraint.
@@ -41,26 +38,22 @@ class SafeSearchConstraintValidator extends ConstraintValidator implements Conta
    * {@inheritdoc}
    */
   public function validate($data, Constraint $constraint) {
-    foreach (\Drupal::service('entity_field.manager')->getFieldDefinitions('node', $data->getType()) as $field_name => $field_def) {
-      // if field is image field.
-      if ($field_def->getType() == 'image') {
-        $settings = $field_def->getThirdPartySettings('google_vision');
-        // if the Safe Search detection is on.
-        if (!empty($settings['google_vision'])) {
-          // if the image is uploaded.
-          if (!empty($data->get($field_name)->target_id)) {
-            // Retrieve the file uri.
-            $file_uri = $data->get($field_name)->entity->getFileUri();
-            if ($filepath = \Drupal::service('file_system')->realpath($file_uri)) {
-              $result = $this->googlevisionapi->safeSearchDetection($filepath);
-              if (!empty($result['responses'][0]['safeSearchAnnotation'])) {
-                $adult = $result['responses'][0]['safeSearchAnnotation']['adult'];
-                $likelihood = array('LIKELY', 'VERY_LIKELY');
-                // if the image has explicit content.
-                if (in_array($adult, $likelihood)) {
-                  $this->context->addViolation($constraint->message, array('%title' => $data->get($field_name)->entity->getFilename()));
-                }
-              }
+    $field_def = $data->getFieldDefinition();
+    $settings = $field_def->getThirdPartySettings('google_vision');
+    // if the Safe Search detection is on.
+    if (!empty($settings['google_vision'])) {
+      // if the image is uploaded.
+      if (!empty($data->getValue('target_id'))) {
+        // Retrieve the file uri.
+        $file_uri = $data->entity->getFileUri();
+        if ($filepath = \Drupal::service('file_system')->realpath($file_uri)) {
+          $result = $this->googlevisionapi->safeSearchDetection($filepath);
+          if (!empty($result['responses'][0]['safeSearchAnnotation'])) {
+            $adult = $result['responses'][0]['safeSearchAnnotation']['adult'];
+            $likelihood = array('LIKELY', 'VERY_LIKELY');
+            // if the image has explicit content.
+            if (in_array($adult, $likelihood)) {
+              $this->context->addViolation($constraint->message);
             }
           }
         }