fix(ui): keep base model selection consistent across training flows

ui/src/Components/BaseModelDropdown.jsx

@@ -1,5 +1,10 @@
 import React from "react";
+import PropTypes from "prop-types";
 import { Dropdown, Option, Field } from "@fluentui/react-components";
+import {
+  applyBaseModelSelection,
+  normalizeBaseModelOptions,
+} from "./BaseModelDropdownHelper";
 
 const styles = {
   container: {
@@ -57,32 +62,21 @@ const renderOption = (option) => {
 function BaseModelDropdown({
   componentState,
   setComponentState,
-  onFormChange,
-
 }) {
   const { baseModelId, baseModelIdError, cataloguedModels = [] } = componentState;
-  const options = React.useMemo(() => {
-    return cataloguedModels.map((m) => ({
-      key: m.key || "none",
-      baseModelName: m.value.baseModelName || "",
-      description: m.value.description.substring(0, 30) + "..." || "",
-      checkpointFilePath: m.value.checkpointFilePath || "",
-      eventTypes: m.value.eventTypes || [],
-      imagerySource: m.value.imagerySource || ""
-    }));
-  }, [cataloguedModels]);
+  const options = React.useMemo(
+    () => normalizeBaseModelOptions(cataloguedModels),
+    [cataloguedModels]
+  );
 
   const selectedOption = options.find(
     (o) => String(o.key) === String(baseModelId)
   );
 
   const handleOptionSelect = (_ev, data) => {
     const picked = options.find((o) => String(o.key) === data.optionValue);
-    onFormChange(
-      picked ? picked.checkpointFilePath : "",
-      "initialWeightsUrl",
-      setComponentState,
-      componentState
+    setComponentState((currentState) =>
+      applyBaseModelSelection(currentState, picked)
     );
   };
 
@@ -106,4 +100,13 @@ function BaseModelDropdown({
   );
 }
 
-export default BaseModelDropdown;
+BaseModelDropdown.propTypes = {
+  componentState: PropTypes.shape({
+    baseModelId: PropTypes.oneOfType([PropTypes.string, PropTypes.number]),
+    baseModelIdError: PropTypes.string,
+    cataloguedModels: PropTypes.array,
+  }).isRequired,
+  setComponentState: PropTypes.func.isRequired,
+};
+
+export default BaseModelDropdown;

ui/src/Components/BaseModelDropdownHelper.js

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+export function buildModelCatalogEndpoint(imageLayer = {}, eventTypes) {
+  const params = new URLSearchParams();
+  const definedEventTypes = Array.isArray(eventTypes)
+    ? eventTypes.filter(Boolean)
+    : [];
+
+  if (definedEventTypes.length > 0) {
+    params.set("eventTypes", definedEventTypes.join(","));
+  }
+
+  const imagerySource = imageLayer?.sourceTypePostEvent;
+  if (imagerySource !== "" && imagerySource != null) {
+    params.set("imagerySource", imagerySource);
+  }
+
+  const query = params.toString();
+  return query ? `GetModelCatalog?${query}` : "GetModelCatalog";
+}
+
+export function buildBaseModelOptionKey(model) {
+  return model?.modelId
+    ? `modelId:${model.modelId}`
+    : `baseModelName:${model?.baseModelName || ""}`;
+}
+
+export function normalizeBaseModelOptions(cataloguedModels = []) {
+  return cataloguedModels.map((model) => {
+    const value = model?.value ?? {};
+    const baseModelName = value.baseModelName || "";
+    const description = value.description == null
+      ? ""
+      : String(value.description);
+
+    return {
+      key: model?.key || `baseModelName:${baseModelName}`,
+      baseModelName,
+      description: description ? `${description.substring(0, 30)}...` : "",
+      checkpointFilePath: value.checkpointFilePath || "",
+      eventTypes: value.eventTypes || [],
+      imagerySource: value.imagerySource || "",
+    };
+  });
+}
+
+export function resolveBaseModelId(cataloguedModels, initialWeightsUrl) {
+  if (!initialWeightsUrl) return "";
+
+  const matchingOption = normalizeBaseModelOptions(cataloguedModels).find(
+    (option) => option.checkpointFilePath === initialWeightsUrl
+  );
+  return matchingOption?.key || "";
+}
+
+export function applyBaseModelSelection(componentState, selectedOption) {
+  return {
+    ...componentState,
+    baseModelId: selectedOption?.key || "",
+    baseModelIdError: "",
+    initialWeightsUrl: selectedOption?.checkpointFilePath || "",
+  };
+}

ui/src/Components/BaseModelDropdownHelper.test.js

@@ -0,0 +1,119 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+
+import {
+  applyBaseModelSelection,
+  buildBaseModelOptionKey,
+  buildModelCatalogEndpoint,
+  normalizeBaseModelOptions,
+  resolveBaseModelId,
+} from "./BaseModelDropdownHelper.js";
+
+test("builds a catalog query with event types and imagery source", () => {
+  const endpoint = buildModelCatalogEndpoint(
+    { sourceTypePostEvent: "Planet" },
+    ["Hurricane", "Flood"]
+  );
+
+  assert.equal(
+    endpoint,
+    "GetModelCatalog?eventTypes=Hurricane%2CFlood&imagerySource=Planet"
+  );
+});
+
+test("omits absent catalog filters instead of stringifying them", () => {
+  assert.equal(buildModelCatalogEndpoint({}, undefined), "GetModelCatalog");
+  assert.equal(
+    buildModelCatalogEndpoint(
+      { sourceTypePostEvent: "Planet" },
+      undefined
+    ),
+    "GetModelCatalog?imagerySource=Planet"
+  );
+});
+
+test("encodes catalog filter values", () => {
+  assert.equal(
+    buildModelCatalogEndpoint(
+      { sourceTypePostEvent: "World View" },
+      ["Severe Storm"]
+    ),
+    "GetModelCatalog?eventTypes=Severe+Storm&imagerySource=World+View"
+  );
+});
+
+test("keeps model ID and fallback name keys in disjoint namespaces", () => {
+  const hasteModelKey = buildBaseModelOptionKey({
+    modelId: "3516",
+    baseModelName: "HASTE model",
+  });
+  const externalModelKey = buildBaseModelOptionKey({
+    baseModelName: "3516",
+  });
+
+  assert.equal(hasteModelKey, "modelId:3516");
+  assert.equal(externalModelKey, "baseModelName:3516");
+  assert.notEqual(hasteModelKey, externalModelKey);
+});
+
+test("normalizes null descriptions and uses model names as fallback keys", () => {
+  const options = normalizeBaseModelOptions([
+    {
+      value: {
+        baseModelName: "External checkpoint A",
+        description: null,
+        checkpointFilePath: "models/external-a.pt",
+      },
+    },
+    {
+      value: {
+        baseModelName: "External checkpoint B",
+        description: null,
+        checkpointFilePath: "models/external-b.pt",
+      },
+    },
+  ]);
+
+  assert.deepEqual(
+    options.map((option) => option.key),
+    [
+      "baseModelName:External checkpoint A",
+      "baseModelName:External checkpoint B",
+    ]
+  );
+  assert.equal(options[0].description, "");
+  assert.equal(options[1].description, "");
+});
+
+test("resolves an existing checkpoint URL to its catalog key", () => {
+  const cataloguedModels = [
+    {
+      key: "modelId:model-1",
+      value: {
+        baseModelName: "Base model",
+        checkpointFilePath: "models/base.pt",
+      },
+    },
+  ];
+
+  assert.equal(
+    resolveBaseModelId(cataloguedModels, "models/base.pt"),
+    "modelId:model-1"
+  );
+  assert.equal(resolveBaseModelId(cataloguedModels, "models/other.pt"), "");
+});
+
+test("applies the selected model id and checkpoint in one state update", () => {
+  const currentState = { name: "Training model", baseModelIdError: "Required" };
+  const selectedOption = {
+    key: "modelId:model-1",
+    checkpointFilePath: "models/base.pt",
+  };
+
+  assert.deepEqual(applyBaseModelSelection(currentState, selectedOption), {
+    name: "Training model",
+    baseModelId: "modelId:model-1",
+    baseModelIdError: "",
+    initialWeightsUrl: "models/base.pt",
+  });
+});

ui/src/Components/CreateEditModelTrainingHelper.js

@@ -2,26 +2,19 @@
 // Licensed under the MIT License.
 
 import { apiGet } from "../util/api";
+import {
+  buildBaseModelOptionKey,
+  buildModelCatalogEndpoint,
+} from "./BaseModelDropdownHelper";
 
 export async function fetchModelCatalog(imageLayer, eventTypes) {
   const cataloguedModels = [];
   try {
-    let eventTypesL = "eventTypes=" + eventTypes;
-
-    if (
-      imageLayer.sourceTypePostEvent !== "" &&
-      imageLayer.sourceTypePostEvent !== null &&
-      imageLayer.sourceTypePostEvent !== undefined
-    ) {
-      const concatChar = eventTypesL === "" ? "" : "&";
-      eventTypesL += concatChar + "imagerySource=" + imageLayer.sourceTypePostEvent;
-    }
-
-    await apiGet(`GetModelCatalog?${eventTypesL}`)
+    await apiGet(buildModelCatalogEndpoint(imageLayer, eventTypes))
       .then((response) => {
         cataloguedModels.push(
           ...response.modelCatalog.map((model) => ({
-            key: model.modelId,
+            key: buildBaseModelOptionKey(model),
             text: model.baseModelName,
             value: model,
           }))
@@ -65,6 +58,8 @@ export function createComponentDefaultState(modelToEdit, imageLayer, projectId)
         batchSizeError: "",
         maxEpochs: "3",
         maxEpochsError: "",
+        baseModelId: "",
+        baseModelIdError: "",
         initialWeightsUrl: "",
         cataloguedModels: [],
         catalogLoading: true,

ui/src/Components/CreateEditModelTrainingModal.jsx

@@ -18,6 +18,7 @@ import { validateEmptyOrInvalid, validateInt, validateFloat } from "../util/vali
 import { useNavigate } from "react-router-dom";
 import { initGuidedTourState, setGuidedTourState } from "./GuidedTourHelper";
 import BaseModelDropdown from "./BaseModelDropdown";
+import { resolveBaseModelId } from "./BaseModelDropdownHelper";
 
 import {
   createComponentDefaultState,
@@ -65,8 +66,14 @@ const CreateEditModelTrainingModal = ({
         projectId
       );
       const cataloguedModels = await fetchModelCatalog(imageLayer, eventTypes);
+      const baseModelId = resolveBaseModelId(
+        cataloguedModels,
+        baseState.initialWeightsUrl
+      );
       setComponentState({
         ...baseState,
+        baseModelId,
+        baseModelIdError: "",
         cataloguedModels,
         catalogLoading: false,
       });
@@ -227,7 +234,6 @@ const CreateEditModelTrainingModal = ({
             <BaseModelDropdown
               componentState={componentState}
               setComponentState={setComponentState}
-              onFormChange={onFormChange}
             />
           </div>
         </div>

ui/src/Components/LabelingTool/LabelingTool.jsx

@@ -48,6 +48,8 @@ const LabelingTool = ({ setModalComponent }) => {
   const [isMapReady, setIsMapReady] = useState(false);
   const [selectedShape, setSelectedShape] = useState(null);
   const [hasUnsavedChanges, setHasUnsavedChanges] = useState(false);
+  const [eventTypes, setEventTypes] = useState([]);
+  const [imageLayer, setImageLayer] = useState(null);
   const { undo, redo } = useDrawingUndoRedo(drawingManager, mapRef);
 
   useEffect(() => {
@@ -181,6 +183,12 @@ const LabelingTool = ({ setModalComponent }) => {
     const projectDetails = await apiGet(
       "GetProjectDetails?projectId=" + projectId
     );
+    setEventTypes(projectDetails.eventTypes || []);
+    setImageLayer(
+      projectDetails.imageLayer?.find(
+        (layer) => layer.imageLayerId === imageLayerId
+      ) || null
+    );
 
     const map = new window.atlas.Map(mapRef.current, {
       center: [0, 0],
@@ -310,6 +318,8 @@ const LabelingTool = ({ setModalComponent }) => {
               setDrawingCount={setDrawingCount}
               selectedShape={selectedShape}
               imageLayerId={imageLayerId}
+              imageLayer={imageLayer}
+              eventTypes={eventTypes}
               undo={undo}
               redo={redo}
             />

ui/src/Components/LabelingTool/LabelingToolRightPanel.jsx

@@ -37,6 +37,8 @@ const LabelingToolRightPanel = ({
   setDrawingCount,
   selectedShape,
   imageLayerId,
+  imageLayer,
+  eventTypes,
   undo,
   redo,
 }) => {
@@ -57,6 +59,8 @@ const LabelingToolRightPanel = ({
     setDrawingCount: PropType.func.isRequired,
     selectedShape: PropType.object,
     imageLayerId: PropType.string.isRequired,
+    imageLayer: PropType.object.isRequired,
+    eventTypes: PropType.array.isRequired,
     undo: PropType.func.isRequired,
     redo: PropType.func.isRequired,
   };
@@ -162,9 +166,10 @@ const LabelingToolRightPanel = ({
           <CreateEditModelTrainingModal
             onClose={() => setModalComponent(null)}
             projectId={projectId}
-            imageLayer={labelingToolDataRef.current}
+            imageLayer={imageLayer}
             guidedTour="createEditModelTrainingModalGuide"
             autoLaunchGuidedTour={true}
+            eventTypes={eventTypes}
           />
         );
       }