From c7a9f97d8e8c3c43361dec0377e3697eba5dc3f4 Mon Sep 17 00:00:00 2001
From: Jake Wagoner <jakewagoneredu@gmail.com>
Date: Mon, 15 Apr 2024 14:47:29 -0600
Subject: [PATCH] Augmentation Pagination and General Error handling (#370)

* Add aug pagination and error handling for deepssm_error

* Add white space to EOF

* Update error handling for deepSSMResult.value

* Update constant name

* Add `-P threads` to the prod celery startup script

---------

Co-authored-by: Anne Haley <anne.haley@kitware.com>
---
 dev/prod.celery.start.sh                      |  2 +-
 .../src/components/DeepSSM/DeepSSMTab.vue     |  6 ++--
 .../src/components/RenderControls.vue         | 34 +++++++++++++++++++
 .../src/components/ShapeViewer/methods.js     | 29 +++++++++-------
 web/shapeworks/src/store/constants.ts         |  2 ++
 web/shapeworks/src/store/index.ts             |  2 ++
 web/shapeworks/src/store/methods.ts           |  2 --
 web/shapeworks/src/views/Main.vue             |  8 +++--
 8 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/dev/prod.celery.start.sh b/dev/prod.celery.start.sh
index 61382db7..33948872 100644
--- a/dev/prod.celery.start.sh
+++ b/dev/prod.celery.start.sh
@@ -5,4 +5,4 @@ cd ~/celery_project
 for keyval in $(cat .env | sed -e 's/: /=/g' -e "s/'\|,\|{\|}//g" -e 's/", "/ /g' -e 's/"}//g' ); do export $keyval; done
 
 conda activate shapeworks
-celery -A shapeworks_cloud.celery worker -n "w1@${HOSTNAME}" -Q gpu --logfile=celery-logs
+celery -A shapeworks_cloud.celery worker -n "w1@${HOSTNAME}" -Q gpu -P threads --logfile=celery-logs
diff --git a/web/shapeworks/src/components/DeepSSM/DeepSSMTab.vue b/web/shapeworks/src/components/DeepSSM/DeepSSMTab.vue
index 142e830a..e9c65c8f 100644
--- a/web/shapeworks/src/components/DeepSSM/DeepSSMTab.vue
+++ b/web/shapeworks/src/components/DeepSSM/DeepSSMTab.vue
@@ -128,9 +128,11 @@ export default {
 
         onMounted(async () => {
             if (!deepSSMResult.value && selectedProject.value) {
+                deepSSMLoadingData.value = true;
                 await loadDeepSSMDataForProject();
+                deepSSMLoadingData.value = false;
             }
-            if (deepSSMResult.value) {
+            if (deepSSMResult.value && deepSSMResult.value.result) {
                 try {
                     Promise.all([
                         await getCSVDataFromURL(deepSSMResult.value.result.aug_total_data),
@@ -311,7 +313,7 @@ export default {
                     <v-btn @click="submitDeepSSMJob">Run DeepSSM tasks</v-btn>
                 </v-expansion-panel-content>
             </v-expansion-panel>
-            <v-expansion-panel v-if="deepSSMResult">
+            <v-expansion-panel v-if="deepSSMResult && deepSSMResult.result">
                 <v-expansion-panel-header>Data</v-expansion-panel-header>
                 <v-expansion-panel-content>
                     <v-tabs v-model="deepSSMDataTab">
diff --git a/web/shapeworks/src/components/RenderControls.vue b/web/shapeworks/src/components/RenderControls.vue
index 8e540b8f..723555cf 100644
--- a/web/shapeworks/src/components/RenderControls.vue
+++ b/web/shapeworks/src/components/RenderControls.vue
@@ -1,6 +1,7 @@
 <script lang="ts">
 import { setDatasetThumbnail, setProjectThumbnail } from '@/api/rest';
 import { computed, ref, watch } from 'vue';
+import _ from 'lodash';
 import {
     particleSize,
     layers,
@@ -29,7 +30,9 @@ import {
     imageViewLevelRange,
     deepSSMResult,
     deepSSMDataTab,
+    deepSSMSamplePage,
     uniformScale,
+    DEEPSSM_SAMPLES_PER_PAGE,
 } from '@/store';
 
 
@@ -181,6 +184,14 @@ export default {
             return false;
         })
 
+        const showSamplePageSelector = computed(() => {
+            return deepSSMDataTab.value === 0 && deepSSMSamplePage.value;
+        })
+
+        const maxSamplePage = computed(() => {
+            return Math.ceil(Object.values(deepSSMResult.value?.aug_pairs).length / DEEPSSM_SAMPLES_PER_PAGE);
+        })
+
         const showUniformScaleOption = computed(() => {
             return deepSSMDataTab.value >= 1; // deepssm data tab training or testing
         })
@@ -232,6 +243,16 @@ export default {
             }
         })
 
+        function changeSamplePage(page: number) {
+            if (page < 1) {
+                page = 1
+            } else if (page > maxSamplePage.value) {
+                page = maxSamplePage.value
+            }
+
+            deepSSMSamplePage.value = page
+        }
+
         return {
             particleSize,
             layersShown,
@@ -259,6 +280,10 @@ export default {
             showAnalysisOptions,
             showUniformScaleOption,
             uniformScale,
+            deepSSMSamplePage,
+            onChangeSamplePage: _.debounce(changeSamplePage, 1000),
+            showSamplePageSelector,
+            maxSamplePage,
             currentTab: props.currentTab,
         }
     }
@@ -313,6 +338,15 @@ export default {
                 v-model="showDifferenceFromMeanMode"
                 label="Show difference from mean"
             />
+            <v-text-field
+                v-if="showSamplePageSelector"
+                :value="deepSSMSamplePage"
+                label="Sample Page"
+                min="0"
+                step="1"
+                @input="onChangeSamplePage"
+                type="number"
+            />
             <v-switch
                 v-if="showUniformScaleOption"
                 v-model="uniformScale"
diff --git a/web/shapeworks/src/components/ShapeViewer/methods.js b/web/shapeworks/src/components/ShapeViewer/methods.js
index 1cc2a891..ba8fce46 100644
--- a/web/shapeworks/src/components/ShapeViewer/methods.js
+++ b/web/shapeworks/src/components/ShapeViewer/methods.js
@@ -250,20 +250,23 @@ export default {
                             }
                             if ([1, 2].includes(deepSSMDataTab.value)) {
                                 const data = shapeData.getPointData().getArrayByName('deepssm_error').getData()
-                                let normalizeRange;
-                                if (uniformScale.value) {
-                                    normalizeRange = deepSSMErrorGlobalRange.value
-                                } else {
-                                    normalizeRange = [Math.min(...data), Math.max(...data)]
+
+                                if (data) {
+                                    let normalizeRange;
+                                    if (uniformScale.value) {
+                                        normalizeRange = deepSSMErrorGlobalRange.value
+                                    } else {
+                                        normalizeRange = [Math.min(...data), Math.max(...data)]
+                                    }
+                                    const normalizedData = data.map((v) => v / (normalizeRange[1] - normalizeRange[0]))
+                                    const normalizedArray = vtkDataArray.newInstance({
+                                        name: 'deepssm_error_normalized',
+                                        values: normalizedData,
+                                    })
+                                    shapeData.getPointData().addArray(normalizedArray)
+                                    mapper.setColorByArrayName('deepssm_error_normalized')
+                                    mapper.setLookupTable(this.lookupTable)
                                 }
-                                const normalizedData = data.map((v) => v / (normalizeRange[1] - normalizeRange[0]))
-                                const normalizedArray = vtkDataArray.newInstance({
-                                    name: 'deepssm_error_normalized',
-                                    values: normalizedData,
-                                })
-                                shapeData.getPointData().addArray(normalizedArray)
-                                mapper.setColorByArrayName('deepssm_error_normalized')
-                                mapper.setLookupTable(this.lookupTable)
                             }
                             const actor = vtkActor.newInstance();
                             if (type) actor.getProperty().setColor(...type.rgb);
diff --git a/web/shapeworks/src/store/constants.ts b/web/shapeworks/src/store/constants.ts
index 69e02404..06a9c66a 100644
--- a/web/shapeworks/src/store/constants.ts
+++ b/web/shapeworks/src/store/constants.ts
@@ -81,3 +81,5 @@ export const COLORS = [
     [255, 255, 153],
     [177, 89, 40],
 ];
+
+export const DEEPSSM_SAMPLES_PER_PAGE = 12;
diff --git a/web/shapeworks/src/store/index.ts b/web/shapeworks/src/store/index.ts
index 839ad043..053bc7f2 100644
--- a/web/shapeworks/src/store/index.ts
+++ b/web/shapeworks/src/store/index.ts
@@ -140,3 +140,5 @@ export const optimizationFormData = ref({});
 export const uniformScale = ref<boolean>(true);
 
 export const deepSSMErrorGlobalRange = ref<number[]>([0, 1]);
+
+export const deepSSMSamplePage = ref<number>(1);
diff --git a/web/shapeworks/src/store/methods.ts b/web/shapeworks/src/store/methods.ts
index 0bb52a75..83ec1002 100644
--- a/web/shapeworks/src/store/methods.ts
+++ b/web/shapeworks/src/store/methods.ts
@@ -163,7 +163,6 @@ export const loadReconstructedSamplesForProject = async (type: string, id: numbe
 
 export const loadDeepSSMDataForProject = async () => {
     if (selectedProject.value) {
-        deepSSMLoadingData.value = true;
         const results = await Promise.all([
             await getDeepSSMResultForProject(
                 selectedProject.value.id
@@ -189,7 +188,6 @@ export const loadDeepSSMDataForProject = async () => {
             images: results[3],
             test_pairs: results[4]
         }
-        deepSSMLoadingData.value = false;
     }
 }
 
diff --git a/web/shapeworks/src/views/Main.vue b/web/shapeworks/src/views/Main.vue
index 6b2a864e..9ed8ce06 100644
--- a/web/shapeworks/src/views/Main.vue
+++ b/web/shapeworks/src/views/Main.vue
@@ -46,7 +46,9 @@ import {
     deepSSMResult,
     deepSSMAugDataShown,
     allDataObjectsInDataset,
-uniformScale,
+    uniformScale,
+    deepSSMSamplePage,
+    DEEPSSM_SAMPLES_PER_PAGE,
 } from '@/store';
 import router from '@/router';
 import TabForm from '@/components/TabForm.vue';
@@ -252,7 +254,8 @@ export default {
                         // populate labelledGroups from aug_pairs
                         // if generated data is shown. Else pass groupedSelections
                         if (deepSSMAugDataShown.value === 'Generated') {
-                            labelledGroups = groupBy(deepSSMResult.value.aug_pairs, 'sample_num')
+                            const paginatedAugPairs = augImages.slice(DEEPSSM_SAMPLES_PER_PAGE * (deepSSMSamplePage.value - 1), DEEPSSM_SAMPLES_PER_PAGE * deepSSMSamplePage.value)
+                            labelledGroups = groupBy(paginatedAugPairs, 'sample_num')
                         }
                         else {
                             labelledGroups = groupedSelections;
@@ -536,6 +539,7 @@ export default {
         watch(deepSSMDataTab, debouncedRefreshRender)
         watch(deepSSMAugDataShown, debouncedRefreshRender)
         watch(uniformScale, debouncedRefreshRender)
+        watch(deepSSMSamplePage, debouncedRefreshRender)
         watch(tab, switchTab)
 
         return {