feat: Migrate annotations table to V2 data (#1433)

* rebase * rebase
chanzuckerberg · Jan 31, 2025 · 1f8a76c · 1f8a76c
1 parent b79eed0
commit 1f8a76c
Show file tree

Hide file tree

Showing 29 changed files with 430 additions and 347 deletions.
diff --git a/frontend/packages/data-portal/app/components/Deposition/MethodLinks/common.tsx b/frontend/packages/data-portal/app/components/Deposition/MethodLinks/common.tsx
@@ -1,8 +1,10 @@
 import { Icon } from '@czi-sds/components'
 import { ReactNode } from 'react'
 
+import { Annotation_Method_Link_Type_Enum } from 'app/__generated_v2__/graphql'
 import { SourceCodeIcon, WeightsIcon } from 'app/components/icons'
 import { VariantLinkProps } from 'app/components/Link'
+import { AnnotationMethodLink } from 'app/types/gql/genericTypes'
 import { I18nKeys } from 'app/types/i18n'
 
 import { METHOD_LINK_TYPES, MethodLinkDataType, MethodLinkType } from './type'
@@ -81,3 +83,26 @@ export function generateMethodLinks(
     )
     .map((props) => methodLinkFromVariant(props))
 }
+
+export function generateMethodLinksV2(
+  links: AnnotationMethodLink[],
+): MethodLinkProps[] {
+  return links
+    .sort(
+      (a, b) =>
+        METHOD_LINK_TYPES.indexOf(getLinkType(a)) -
+        METHOD_LINK_TYPES.indexOf(getLinkType(b)),
+    )
+    .map((link) => ({
+      title: link.name ?? '',
+      url: link.link ?? '',
+      icon: ICON_MAP[getLinkType(link)],
+      i18nLabel: METHOD_TYPE_TO_I18N_KEY[getLinkType(link)],
+    }))
+}
+
+function getLinkType(
+  link: AnnotationMethodLink,
+): Annotation_Method_Link_Type_Enum {
+  return link.linkType ?? Annotation_Method_Link_Type_Enum.Other
+}
diff --git a/frontend/packages/data-portal/app/components/Download/ConfigureAnnotationDownloadContent.tsx b/frontend/packages/data-portal/app/components/Download/ConfigureAnnotationDownloadContent.tsx
@@ -1,51 +1,34 @@
-import { useMemo } from 'react'
-import { DeepPartial } from 'utility-types'
-
-import { AnnotationFileEdge } from 'app/__generated_v2__/graphql'
 import { useDownloadModalContext } from 'app/context/DownloadModal.context'
 import { useDownloadModalQueryParamState } from 'app/hooks/useDownloadModalQueryParamState'
-import { useRunById } from 'app/hooks/useRunById'
 
 import { AnnotationAlignmentCallout } from './AnnotationAlignmentCallout'
 import { FileFormatDropdown } from './FileFormatDropdown'
 
 export function ConfigureAnnotationDownloadContent() {
-  const { objectShapeType, fileFormat } = useDownloadModalQueryParamState()
-  const { annotationToDownload, allTomograms } = useDownloadModalContext()
-
-  const fileFormats = useMemo<string[]>(
-    () =>
-      annotationToDownload?.files
-        .filter((annotation) => annotation.shape_type === objectShapeType)
-        .map((annotation) => annotation.format) ?? [],
-    [annotationToDownload?.files, objectShapeType],
-  )
-
-  const { annotationShapes } = useRunById()
-
-  const isMatchingFormat = (file: DeepPartial<AnnotationFileEdge>) =>
-    file?.node?.format === fileFormat
-
-  const alignmentId =
-    annotationShapes
-      .find(
-        (shape) =>
-          shape.annotation?.id === annotationToDownload?.id &&
-          shape.shapeType === objectShapeType &&
-          shape.annotationFiles.edges.some(isMatchingFormat),
-      )
-      ?.annotationFiles.edges.find(isMatchingFormat)?.node.alignmentId ?? 0
+  const { fileFormat } = useDownloadModalQueryParamState()
+  const { annotationShapeToDownload, allTomograms } = useDownloadModalContext()
+
+  const fileFormats =
+    annotationShapeToDownload?.annotationFiles.edges.map(
+      (file) => file.node.format,
+    ) ?? []
+  const annotationFileAlignmentId =
+    annotationShapeToDownload?.annotationFiles.edges.find(
+      (file) => file.node.format === fileFormat,
+    )?.node.alignmentId ?? undefined
 
   return (
     <>
       <FileFormatDropdown className="pt-sds-l" fileFormats={fileFormats} />
-      <AnnotationAlignmentCallout
-        alignmentId={alignmentId}
-        initialState="open"
-        misalignedTomograms={(allTomograms ?? []).filter(
-          (tomogram) => tomogram.alignment?.id !== alignmentId,
-        )}
-      />
+      {annotationFileAlignmentId !== undefined && (
+        <AnnotationAlignmentCallout
+          alignmentId={annotationFileAlignmentId}
+          initialState="open"
+          misalignedTomograms={(allTomograms ?? []).filter(
+            (tomogram) => tomogram.alignment?.id !== annotationFileAlignmentId,
+          )}
+        />
+      )}
     </>
   )
 }
diff --git a/frontend/packages/data-portal/app/components/Download/ConfigureDownloadContent.tsx b/frontend/packages/data-portal/app/components/Download/ConfigureDownloadContent.tsx
@@ -9,13 +9,16 @@ import { ConfigureTomogramDownloadContent } from './ConfigureTomogramDownloadCon
 
 export function ConfigureDownloadContent() {
   const { t } = useI18n()
-
-  const { annotationName, annotationId, objectShapeType } =
+  const { annotationName, annotationId, fileFormat, objectShapeType } =
     useDownloadModalQueryParamState()
-
-  const { annotationToDownload, datasetTitle, runName, objectName } =
+  const { annotationShapeToDownload, datasetTitle, runName, objectName } =
     useDownloadModalContext()
 
+  const annotationFileAlignmentId =
+    annotationShapeToDownload?.annotationFiles.edges.find(
+      (file) => file.node.format === fileFormat,
+    )?.node.alignmentId ?? undefined
+
   return (
     <>
       <ModalSubtitle label={t('datasetName')} value={datasetTitle} />
@@ -35,10 +38,10 @@ export function ConfigureDownloadContent() {
       {objectShapeType && (
         <ModalSubtitle label={t('objectShapeType')} value={objectShapeType} />
       )}
-      {annotationToDownload !== undefined && (
+      {annotationFileAlignmentId !== undefined && (
         <ModalSubtitle
           label={t('alignmentId')}
-          value={`${IdPrefix.Alignment}-${annotationToDownload.id}`}
+          value={`${IdPrefix.Alignment}-${annotationFileAlignmentId}`}
         />
       )}
 

diff --git a/frontend/packages/data-portal/app/components/Download/ConfigureTomogramDownloadContent.tsx b/frontend/packages/data-portal/app/components/Download/ConfigureTomogramDownloadContent.tsx
@@ -29,7 +29,7 @@ export function ConfigureTomogramDownloadContent() {
 
   const {
     tomogramToDownload,
-    allAnnotationFiles = [],
+    allAnnotationShapes = [],
     allTomograms = [],
     runId,
   } = useDownloadModalContext()
@@ -84,7 +84,7 @@ export function ConfigureTomogramDownloadContent() {
         <Radio
           value={DownloadConfig.AllAnnotations}
           label={t('downloadAllAnnotations')}
-          disabled={allAnnotationFiles.length === 0}
+          disabled={allAnnotationShapes.length === 0}
           disabledTooltip={t('noAnnotationsAvailableToDownload')}
           description={t('downloadAvailableAnnotationsInSupported')}
           onClick={setAllAnnotationsConfig}

diff --git a/frontend/packages/data-portal/app/components/Download/DownloadOptionsContent.tsx b/frontend/packages/data-portal/app/components/Download/DownloadOptionsContent.tsx
@@ -54,7 +54,7 @@ export function DownloadOptionsContent() {
     objectName,
     runId,
     runName,
-    annotationToDownload,
+    annotationShapeToDownload,
     tomogramToDownload,
     type,
   } = useDownloadModalContext()
@@ -65,6 +65,10 @@ export function DownloadOptionsContent() {
   const referenceTomogram = allTomograms?.find(
     (tomogram) => tomogram.id.toString() === referenceTomogramId,
   )
+  const annotationFileAlignmentId =
+    annotationShapeToDownload?.annotationFiles.edges.find(
+      (file) => file.node.format === fileFormat,
+    )?.node.alignmentId ?? undefined
   const DownloadTabContent = DOWNLOAD_TAB_MAP[selectedTab]
 
   return (
@@ -115,16 +119,10 @@ export function DownloadOptionsContent() {
           value={getTomogramName(referenceTomogram)}
         />
       )}
-      {annotationToDownload && tomogramToDownload?.alignment && (
+      {annotationFileAlignmentId !== undefined && (
         <ModalSubtitle
           label={t('alignmentId')}
-          value={tomogramToDownload.alignment.id}
-        />
-      )}
-      {annotationToDownload && (
-        <ModalSubtitle
-          label={t('alignmentId')}
-          value={`${IdPrefix.Alignment}-${annotationToDownload.id}`}
+          value={`${IdPrefix.Alignment}-${annotationFileAlignmentId}`}
         />
       )}
       {fileFormat && (
@@ -163,11 +161,16 @@ export function DownloadOptionsContent() {
 
       <DownloadTabContent />
 
-      {annotationToDownload !== undefined && tomogramToDownload?.alignment ? (
+      {annotationFileAlignmentId !== undefined ? (
         <AnnotationAlignmentCallout
-          alignmentId={tomogramToDownload.alignment.id}
+          alignmentId={annotationFileAlignmentId}
           initialState="closed"
-          misalignedTomograms={[]}
+          misalignedTomograms={
+            allTomograms?.filter(
+              (tomogram) =>
+                tomogram.alignment?.id !== annotationFileAlignmentId,
+            ) ?? []
+          }
         />
       ) : (
         <Callout intent="notice" className="!w-full !mt-sds-xl">

diff --git a/frontend/packages/data-portal/app/components/Download/FileFormatDropdown.tsx b/frontend/packages/data-portal/app/components/Download/FileFormatDropdown.tsx
@@ -21,6 +21,17 @@ const FILE_FORMAT_TOOLTIP_I18N: Record<string, I18nKeys> = {
 }
 
 export const FILE_FORMAT_ORDER = ['mrc', 'zarr', 'ndjson']
+export function getDefaultFileFormat(
+  availableFormats: string[],
+): string | undefined {
+  for (const fileFormat of FILE_FORMAT_ORDER) {
+    if (availableFormats.includes(fileFormat)) {
+      return fileFormat
+    }
+  }
+
+  return undefined
+}
 
 /**
  * Renders select dropdown with file formats specified in the `fileFormats`
@@ -41,7 +52,6 @@ export function FileFormatDropdown({
   const matchingFileFormats = FILE_FORMAT_ORDER.filter((format) =>
     fileFormats.includes(format),
   )
-  const defaultFormat = matchingFileFormats[0]
 
   const fileFormatOptions = useMemo<SelectOption[]>(
     () =>
@@ -53,7 +63,8 @@ export function FileFormatDropdown({
     [matchingFileFormats, t],
   )
 
-  const selectedFormat = fileFormat ?? defaultFormat
+  const selectedFormat =
+    fileFormat ?? getDefaultFileFormat(matchingFileFormats)!
 
   return (
     <Select

diff --git a/frontend/packages/data-portal/app/components/Filters/AnnotationObjectShapeTypeFilter.tsx b/frontend/packages/data-portal/app/components/Filters/AnnotationObjectShapeTypeFilter.tsx
@@ -1,28 +1,28 @@
 import { useMemo } from 'react'
 
-import { shapeTypeToI18nKey } from 'app/constants/objectShapeTypes'
+import { Annotation_File_Shape_Type_Enum } from 'app/__generated_v2__/graphql'
+import { getShapeTypeI18nKey } from 'app/constants/objectShapeTypes'
 import { QueryParams } from 'app/constants/query'
 import { useFilter } from 'app/hooks/useFilter'
 import { useI18n } from 'app/hooks/useI18n'
 import { BaseFilterOption } from 'app/types/filter'
-import { ObjectShapeType } from 'app/types/shapeTypes'
 
 import { SelectFilter } from './SelectFilter'
 
 function getObjectShapeTypeOptions(
-  objectShapeTypes: string[],
+  objectShapeTypes: Annotation_File_Shape_Type_Enum[],
   t: ReturnType<typeof useI18n>['t'],
 ): BaseFilterOption[] {
   return objectShapeTypes.map((value) => ({
-    label: t(shapeTypeToI18nKey[value as ObjectShapeType]),
+    label: t(getShapeTypeI18nKey(value)),
     value,
   }))
 }
 
 export function AnnotatedObjectShapeTypeFilter({
   allObjectShapeTypes,
 }: {
-  allObjectShapeTypes: string[]
+  allObjectShapeTypes: Annotation_File_Shape_Type_Enum[]
 }) {
   const { t } = useI18n()
 

diff --git a/frontend/packages/data-portal/app/components/Run/AnnotationConfidenceTable.tsx b/frontend/packages/data-portal/app/components/Run/AnnotationConfidenceTable.tsx
@@ -1,16 +1,16 @@
 import { AccordionMetadataTable } from 'app/components/AccordionMetadataTable'
 import { useI18n } from 'app/hooks/useI18n'
-import { useAnnotation } from 'app/state/annotation'
+import { useSelectedAnnotationShape } from 'app/state/annotation'
 
 export function AnnotationConfidenceTable() {
-  const { activeAnnotation: annotation } = useAnnotation()
+  const { selectedAnnotationShape } = useSelectedAnnotationShape()
   const { t } = useI18n()
 
-  if (!annotation) {
+  if (!selectedAnnotationShape) {
     return null
   }
 
-  const isGroundTruth = annotation.ground_truth_status
+  const isGroundTruth = selectedAnnotationShape.annotation?.groundTruthStatus
 
   return (
     <AccordionMetadataTable
@@ -26,7 +26,7 @@ export function AnnotationConfidenceTable() {
           values: [
             isGroundTruth
               ? t('notApplicable')
-              : annotation.ground_truth_used ?? '--',
+              : selectedAnnotationShape.annotation?.groundTruthUsed ?? '--',
           ],
           className: 'text-sds-color-primitive-gray-500',
         },
@@ -35,7 +35,7 @@ export function AnnotationConfidenceTable() {
           values: [
             isGroundTruth
               ? t('notApplicable')
-              : annotation.confidence_precision ?? '--',
+              : selectedAnnotationShape.annotation?.confidenceRecall ?? '--',
           ],
           className: 'text-sds-color-primitive-gray-500',
         },
@@ -44,7 +44,7 @@ export function AnnotationConfidenceTable() {
           values: [
             isGroundTruth
               ? t('notApplicable')
-              : annotation.confidence_recall ?? '--',
+              : selectedAnnotationShape.annotation?.confidenceRecall ?? '--',
           ],
           className: 'text-sds-color-primitive-gray-500',
         },

diff --git a/frontend/packages/data-portal/app/components/Run/AnnotationDrawer.tsx b/frontend/packages/data-portal/app/components/Run/AnnotationDrawer.tsx
@@ -2,27 +2,27 @@ import { MetadataDrawer } from 'app/components/MetadataDrawer'
 import { IdPrefix } from 'app/constants/idPrefixes'
 import { MetadataDrawerId } from 'app/hooks/useMetadataDrawer'
 import { i18n } from 'app/i18n'
-import { useAnnotation } from 'app/state/annotation'
-import { getAnnotationTitle } from 'app/utils/annotation'
+import { useSelectedAnnotationShape } from 'app/state/annotation'
 
 import { AnnotationConfidenceTable } from './AnnotationConfidenceTable'
 import { AnnotationObjectTable } from './AnnotationObjectTable/AnnotationObjectTable'
 import { AnnotationOverviewTable } from './AnnotationOveriewTable'
 
 export function AnnotationDrawer() {
-  const { activeAnnotation, setActiveAnnotation } = useAnnotation()
+  const { selectedAnnotationShape, setSelectedAnnotationShape } =
+    useSelectedAnnotationShape()
 
   return (
     <MetadataDrawer
-      disabled={!activeAnnotation}
+      disabled={!selectedAnnotationShape}
       drawerId={MetadataDrawerId.Annotation}
       label={i18n.annotationDetails}
-      title={getAnnotationTitle(activeAnnotation)}
+      title={`${selectedAnnotationShape?.id} ${selectedAnnotationShape?.annotation?.objectName}`}
       idInfo={{
         label: 'annotationId',
-        text: `${IdPrefix.Annotation}-${activeAnnotation?.id}`,
+        text: `${IdPrefix.Annotation}-${selectedAnnotationShape?.annotation?.id}`,
       }}
-      onClose={() => setActiveAnnotation(null)}
+      onClose={() => setSelectedAnnotationShape(null)}
     >
       <AnnotationOverviewTable />
       <AnnotationObjectTable />