feat: add layout analysis stack

awslabs · Mar 5, 2024 · e3925e0 · e3925e0
1 parent dd592a4
commit e3925e0
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 7 deletions.
diff --git a/src/api-deployment/ai-solution-kit-stack.ts b/src/api-deployment/ai-solution-kit-stack.ts
@@ -63,6 +63,7 @@ import { ObjectRecognitionSageMakerFeatureNestedStack } from './features/sagemak
 import { PornographyDetectionSageMakerFeatureNestedStack } from './features/sagemaker/pornography-detection-sagemaker';
 import { SuperResolutionSageMakerFeatureNestedStack } from './features/sagemaker/super-resolution-sagemaker';
 import { TextSimilaritySageMakerFeatureNestedStack } from './features/sagemaker/text-similarity-sagemaker';
+import { LayoutAnalysisSageMakerFeatureNestedStack } from './features/sagemaker/layout-analysis-sagemaker';
 
 export interface FeatureProps {
   readonly featureStack: FeatureNestedStack;
@@ -521,6 +522,19 @@ export class AISolutionKitStack extends Stack {
       this.addOutput(cfnTemplate, api.restApiId, 'general-nlu-ml', 'General NLU SageMaker', 'ConditionGeneralNLUSageMaker');
     }
 
+    // Feature: Layout Analysis SageMaker
+    {
+      const layoutAnalysisSageMaker = new LayoutAnalysisSageMakerFeatureNestedStack(this, 'Layout-Analysis-SageMaker', {
+        restApi: api,
+        customAuthorizationType: authType,
+        ecrDeployment: ecrDeployment,
+        updateCustomResourceProvider: updateCustomResourceProvider,
+        ecrRegistry: props.ecrRegistry,
+      });
+      (layoutAnalysisSageMaker.nestedStackResource as CfnStack).cfnOptions.condition = cfnTemplate.getCondition('ConditionLayoutAnalysisSageMaker');
+      this.addOutput(cfnTemplate, api.restApiId, 'layout-analysis-ml', 'Layout Analysis SageMaker', 'ConditionLayoutAnalysisSageMaker');
+    }
+
     // Stage base URL
     {
       new CfnOutput(this, 'Stage base URL', {

diff --git a/src/api-deployment/features/api-explorer/openapi.yaml b/src/api-deployment/features/api-explorer/openapi.yaml
@@ -124,6 +124,40 @@ paths:
           description: Internal Server Error
         '504':
           description: Endpoint request timed out
+  /layout-analysis:
+    post:
+      tags:
+        - Optical Character Recognition(OCR)
+      summary: Convert document scan images to a Markdown format document.
+      operationId: layout-analysis
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/imageRequest'
+            examples:
+              URL:
+                $ref: '#/components/examples/imageURL'
+              Base64:
+                $ref: '#/components/examples/imageBase64'
+      responses:
+        '200':
+          description: successful operation
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  $ref: '#/components/schemas/ocrResponse'
+        '403':
+          description: Missing Authentication Token
+        '500':
+          description: Internal server error
+        '502':
+          description: Internal Server Error
+        '504':
+          description: Endpoint request timed out
   /custom-ocr:
     post:
       tags:

diff --git a/src/api-deployment/features/sagemaker/layout-analysis-sagemaker.ts b/src/api-deployment/features/sagemaker/layout-analysis-sagemaker.ts
@@ -0,0 +1,26 @@
+import { CustomResource, RemovalPolicy } from 'aws-cdk-lib';
+import { Repository } from 'aws-cdk-lib/aws-ecr';
+import { Construct } from 'constructs';
+import { FeatureNestedStack, FeatureNestedStackProps } from '../../feature-nested-stack';
+import { SageMakerFeatureConstruct } from '../../sagemaker-feature-construct';
+
+export class LayoutAnalysisSageMakerFeatureNestedStack extends FeatureNestedStack {
+  constructor(scope: Construct, id: string, props: FeatureNestedStackProps) {
+
+    super(scope, id, props);
+    const featureName = 'layout-analysis';
+    this.templateOptions.description = '(SO8023-lag) - AI Solution Kit - Layout Analysis. Template version v1.4.0. See https://awslabs.github.io/aws-ai-solution-kit/en/layout-analysis.';
+
+    // The default instance type is ml.g4dn.xlarge
+    const sageMakerConstruct = new SageMakerFeatureConstruct(this, `${featureName}-construct`, {
+      rootRestApi: props.restApi,
+      authorizationType: props.customAuthorizationType,
+      restApiResourcePath: `${featureName}`,
+      featureName: `${featureName}`,
+      featureCategory: 'media',
+      updateCustomResourceProvider: props.updateCustomResourceProvider,
+      sageMakerInstanceType: props.instanceType == null ? 'ml.g4dn.xlarge' : props.instanceType,
+      sageMakerEcrDeployment: props.ecrDeployment,
+    });
+  }
+}
diff --git a/src/api-deployment/parameter-group.template b/src/api-deployment/parameter-group.template
@@ -15,6 +15,7 @@ Metadata:
           - GeneralOCR
           - GeneralOCRTraditionalChinese
           - AdvancedOCRSageMaker
+          - LayoutAnalysis
           - CustomOCR
           - CarLicensePlate
 
@@ -251,6 +252,14 @@ Parameters:
       - "yes-sagemaker"
       - "no"
 
+  LayoutAnalysis:
+    Description: Convert document scan images to a Markdown format document. Choose yes to activate the API. See https://awslabs.github.io/aws-ai-solution-kit/en/deploy-layout-analysis.
+    Default: "no"
+    Type: String
+    AllowedValues:
+      - "yes-sagemaker"
+      - "no"
+
 Conditions:
   ConditionAPIExplorer: !Equals 
       - !Ref APIExplorer
@@ -366,4 +375,8 @@ Conditions:
 
   ConditionGeneralNLUSageMaker: !Equals
       - !Ref GeneralNLU
+      - "yes-sagemaker"
+
+  ConditionLayoutAnalysisSageMaker: !Equals
+      - !Ref LayoutAnalysis
       - "yes-sagemaker"
diff --git a/src/main.ts b/src/main.ts
@@ -32,13 +32,13 @@ if (buildContainers === 'true' || deployContainers === 'true') {
       app: 'ai-solution-kit',
     },
   });
-  new AISolutionKitWorkshopStack(app, 'AI-Solution-Kit-Workshop', {
-    synthesizer: synthesizer(),
-    ecrRegistry: ecrRegistry === 'undefined' ? 'public.ecr.aws/aws-gcr-solutions/aws-gcr-ai-solution-kit' : ecrRegistry,
-    tags: {
-      app: 'ai-solution-kit',
-    },
-  });
+  // new AISolutionKitWorkshopStack(app, 'AI-Solution-Kit-Workshop', {
+  //   synthesizer: synthesizer(),
+  //   ecrRegistry: ecrRegistry === 'undefined' ? 'public.ecr.aws/aws-gcr-solutions/aws-gcr-ai-solution-kit' : ecrRegistry,
+  //   tags: {
+  //     app: 'ai-solution-kit',
+  //   },
+  // });
 }
 
 if (process.env.USE_BSS) {