Skip to content

Commit 43332a4

Browse files
hotzenklotzhotzenklotz
committed
added alignment job AI UI
1 parent d1cacc6 commit 43332a4

16 files changed

+572
-93
lines changed

‎frontend/javascripts/libs/toast.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,8 @@ const Toast = {
9595
<Collapse
9696
className="collapsibleToastDetails"
9797
bordered={false}
98+
ghost
9899
style={{
99-
background: "transparent",
100100
marginLeft: -16,
101101
}}
102102
items={[

‎frontend/javascripts/viewer/view/action-bar/ai_job_modals/components/collapsible_split_merger_evaluation_settings.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export function CollapsibleSplitMergerEvaluationSettings({
1919
style={{ marginBottom: 8 }}
2020
onChange={handleCollapseChange}
2121
expandIcon={() => <Checkbox checked={isActive} />}
22+
ghost
2223
items={[
2324
{
2425
key: "evaluation",

‎frontend/javascripts/viewer/view/action-bar/ai_job_modals/components/should_use_trees_form_item.tsx‎

Lines changed: 22 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import type { RuleObject } from "antd/es/form";
44
import { useWkSelector } from "libs/react_hooks";
55
import { useCallback, useMemo } from "react";
66

7-
export function ShouldUseTreesFormItem() {
7+
export function ShouldUseManualMatchesFormItem() {
88
const annotation = useWkSelector((state) => state.annotation);
99
const trees = useMemo(
1010
() => (annotation.skeleton ? annotation.skeleton.trees.values().toArray() : []),
@@ -36,28 +36,26 @@ export function ShouldUseTreesFormItem() {
3636
);
3737

3838
return (
39-
<div>
40-
<Form.Item
41-
name="useAnnotation"
42-
label={
43-
<Space>
44-
<div style={{}}>
45-
Manual Matches{" "}
46-
<Tooltip title="Please select whether the alignment should take connected skeleton nodes between adjacent sections as alignment guideline whenever available.">
47-
<InfoCircleOutlined />
48-
</Tooltip>
49-
</div>
50-
</Space>
51-
}
52-
valuePropName="checked"
53-
rules={[
54-
{
55-
validator,
56-
},
57-
]}
58-
>
59-
<Checkbox> Use manual matches from skeleton. </Checkbox>
60-
</Form.Item>
61-
</div>
39+
<Form.Item
40+
name="useAnnotation"
41+
label={
42+
<Space>
43+
<div style={{}}>
44+
Manual Matches{" "}
45+
<Tooltip title="Please select whether the alignment should take connected skeleton nodes between adjacent sections as alignment guideline whenever available.">
46+
<InfoCircleOutlined />
47+
</Tooltip>
48+
</div>
49+
</Space>
50+
}
51+
valuePropName="checked"
52+
rules={[
53+
{
54+
validator,
55+
},
56+
]}
57+
>
58+
<Checkbox> Use manual matches from skeleton. </Checkbox>
59+
</Form.Item>
6260
);
6361
}

‎frontend/javascripts/viewer/view/action-bar/ai_job_modals/forms/start_job_form.tsx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import type { UserBoundingBox } from "viewer/store";
2121
import { BoundingBoxSelectionFormItem } from "../components/bounding_box_selection_form_item";
2222
import { CollapsibleWorkflowYamlEditor } from "../components/collapsible_workflow_yaml_editor";
2323
import { JobCreditCostInformation } from "../components/job_credit_cost_information";
24-
import { ShouldUseTreesFormItem } from "../components/should_use_trees_form_item";
24+
import { ShouldUseManualMatchesFormItem } from "../components/should_use_trees_form_item";
2525
import { useCurrentlySelectedBoundingBox } from "../hooks/use_currently_selected_bounding_box";
2626
import DEFAULT_PREDICT_WORKFLOW from "../templates/default-predict-workflow-template";
2727
import { getBoundingBoxesForLayers } from "../utils";
@@ -239,7 +239,7 @@ export function StartJobForm(props: StartJobFormProps) {
239239
showVolume={jobCreditCostPerGVx != null}
240240
/>
241241
{jobSpecificInputFields}
242-
{isSkeletonSelectable && <ShouldUseTreesFormItem />}
242+
{isSkeletonSelectable && <ShouldUseManualMatchesFormItem />}
243243
{props.showWorkflowYaml ? (
244244
<CollapsibleWorkflowYamlEditor
245245
isActive={useCustomWorkflow}

‎frontend/javascripts/viewer/view/ai_jobs/ai_jobs_drawer.tsx‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,22 @@
11
import { Drawer, Tabs } from "antd";
2-
import { AiImageAlignment } from "./alignment/ai_image_alignment_job";
3-
import { AiImageSegmentation } from "./run_ai_model/ai_image_segmentation_job";
2+
import { AiImageAlignmentJob } from "./alignment/ai_image_alignment_job";
3+
import { AiImageSegmentationJob } from "./run_ai_model/ai_image_segmentation_job";
44
import { AiModelTraining } from "./train_ai_model/ai_model_training_job";
55

66
const { TabPane } = Tabs;
77

88
export const AiJobsDrawer = () => {
99
return (
1010
<Drawer title="AI Jobs" placement="right" width={1200} open={true}>
11-
<Tabs defaultActiveKey="1">
11+
<Tabs defaultActiveKey="3">
1212
<TabPane tab="Image Segmentation" key="1">
13-
<AiImageSegmentation />
13+
<AiImageSegmentationJob />
1414
</TabPane>
1515
<TabPane tab="Model Training" key="2">
1616
<AiModelTraining />
1717
</TabPane>
1818
<TabPane tab="Image Alignment" key="3">
19-
<AiImageAlignment />
19+
<AiImageAlignmentJob />
2020
</TabPane>
2121
</Tabs>
2222
</Drawer>

‎frontend/javascripts/viewer/view/ai_jobs/alignment/ai_alignment_job_context.tsx‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import { startAlignSectionsJob } from "admin/rest_api";
2+
import { useWkSelector } from "libs/react_hooks";
3+
import Toast from "libs/toast";
4+
import type React from "react";
5+
import { createContext, useCallback, useContext, useEffect, useState } from "react";
6+
import { useDispatch } from "react-redux";
7+
import type { APIJobType } from "types/api_types";
8+
import { getColorLayers } from "viewer/model/accessors/dataset_accessor";
9+
import { setAIJobModalStateAction } from "viewer/model/actions/ui_actions";
10+
import type { UserBoundingBox } from "viewer/store";
11+
import { getBoundingBoxesForLayers } from "viewer/view/action-bar/ai_job_modals/utils";
12+
import type { AlignmentTask } from "./ai_alignment_model_selector";
13+
14+
interface AlignmentJobContextType {
15+
handleStartAnalysis: () => void;
16+
newDatasetName: string;
17+
selectedTask: AlignmentTask | null;
18+
selectedJobType: APIJobType | null;
19+
setSelectedJobType: (jobType: APIJobType) => void;
20+
setSelectedTask: (task: AlignmentTask) => void;
21+
selectedBoundingBox: UserBoundingBox | null;
22+
setNewDatasetName: (name: string) => void;
23+
shouldUseManualMatches: boolean;
24+
setShouldUseManualMatches: (shouldUseManualMatches: boolean) => void;
25+
}
26+
27+
const AlignmentJobContext = createContext<AlignmentJobContextType | undefined>(undefined);
28+
29+
export const AlignmentJobContextProvider: React.FC<{ children: React.ReactNode }> = ({
30+
children,
31+
}) => {
32+
const [selectedTask, setSelectedTask] = useState<AlignmentTask | null>(null);
33+
const [selectedJobType, setSelectedJobType] = useState<APIJobType | null>(null);
34+
const [newDatasetName, setNewDatasetName] = useState("");
35+
const [shouldUseManualMatches, setShouldUseManualMatches] = useState(false);
36+
const dispatch = useDispatch();
37+
38+
const dataset = useWkSelector((state) => state.dataset);
39+
const annotationId = useWkSelector((state) => state.annotation.annotationId);
40+
const colorLayers = getColorLayers(dataset);
41+
const colorLayer = colorLayers[0];
42+
43+
const selectedBoundingBox = getBoundingBoxesForLayers([colorLayer])[0];
44+
45+
useEffect(() => {
46+
if (dataset && selectedTask) {
47+
setNewDatasetName(`${dataset.name}_${selectedTask.name?.replace(/\s/g, "_")}`);
48+
}
49+
}, [dataset, selectedTask]);
50+
51+
const handleStartAnalysis = useCallback(async () => {
52+
try {
53+
startAlignSectionsJob(
54+
dataset.id,
55+
colorLayer.name,
56+
newDatasetName,
57+
shouldUseManualMatches ? annotationId : undefined,
58+
);
59+
Toast.success("Alignment started successfully!");
60+
dispatch(setAIJobModalStateAction("invisible"));
61+
} catch (error) {
62+
console.error(error);
63+
Toast.error("Failed to start alignment.");
64+
}
65+
}, [dataset.id, dispatch, colorLayer.name, newDatasetName, annotationId, shouldUseManualMatches]);
66+
67+
const value = {
68+
selectedJobType,
69+
selectedTask,
70+
selectedBoundingBox,
71+
newDatasetName,
72+
setSelectedJobType,
73+
setNewDatasetName,
74+
setSelectedTask,
75+
handleStartAnalysis,
76+
shouldUseManualMatches,
77+
setShouldUseManualMatches,
78+
};
79+
80+
return <AlignmentJobContext.Provider value={value}>{children}</AlignmentJobContext.Provider>;
81+
};
82+
83+
export const useAlignmentJobContext = () => {
84+
const context = useContext(AlignmentJobContext);
85+
if (context === undefined) {
86+
throw new Error("useAlignmentJobContext must be used within a AlignmentJobContextProvider");
87+
}
88+
return context;
89+
};

‎frontend/javascripts/viewer/view/ai_jobs/alignment/ai_alignment_model_selector.tsx‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import { ExperimentOutlined } from "@ant-design/icons";
2+
import { Avatar, Card, List, Space, Tag, Typography } from "antd";
3+
import type React from "react";
4+
import { useCallback } from "react";
5+
import { APIJobType } from "types/api_types";
6+
import { useAlignmentJobContext } from "./ai_alignment_job_context";
7+
8+
const { Text } = Typography;
9+
10+
export type AlignmentTask = {
11+
name: string;
12+
comment: string;
13+
id: string;
14+
jobType: APIJobType | null;
15+
image: string;
16+
disabled?: boolean;
17+
};
18+
19+
const alignmentTasks: AlignmentTask[] = [
20+
{
21+
name: "Align Sections",
22+
comment:
23+
"Align all sections of this dataset along the Z axis using features in neighboring sections.",
24+
id: "align-sections",
25+
jobType: APIJobType.ALIGN_SECTIONS,
26+
image: "/assets/images/align_example.png",
27+
},
28+
{
29+
name: "Align multiple tiles",
30+
comment: "For aligning datasets with multiple tiles per section, please contact us.",
31+
id: "align-tiles",
32+
disabled: true,
33+
jobType: null,
34+
image: "/assets/images/align_example.png",
35+
},
36+
];
37+
38+
export const AiAlignmentModelSelector: React.FC = () => {
39+
const { setSelectedJobType, selectedTask, setSelectedTask } = useAlignmentJobContext();
40+
41+
const handleTaskSelection = useCallback(
42+
(item: AlignmentTask) => {
43+
if (!item.disabled && item.jobType) {
44+
setSelectedTask(item);
45+
setSelectedJobType(item.jobType);
46+
}
47+
},
48+
[setSelectedJobType, setSelectedTask],
49+
);
50+
51+
return (
52+
<Card
53+
title={
54+
<Space align="center">
55+
<ExperimentOutlined style={{ color: "#1890ff" }} />
56+
Select AI Alignment Task
57+
</Space>
58+
}
59+
>
60+
<List
61+
itemLayout="horizontal"
62+
dataSource={alignmentTasks}
63+
renderItem={(item) => (
64+
<List.Item
65+
style={{
66+
border: selectedTask?.id === item.id ? "1px solid #1890ff" : "1px solid #d9d9d9",
67+
borderRadius: "8px",
68+
marginBottom: "16px",
69+
padding: "16px",
70+
opacity: item.disabled ? 0.5 : 1,
71+
cursor: item.disabled ? "not-allowed" : "pointer",
72+
}}
73+
onClick={() => handleTaskSelection(item)}
74+
>
75+
<List.Item.Meta
76+
avatar={
77+
<Avatar shape="square" size={64} src={<img src={item.image} alt={item.name} />} />
78+
}
79+
title={
80+
<Space>
81+
<Text strong>{item.name}</Text>
82+
{item.disabled && <Tag>Coming Soon</Tag>}
83+
</Space>
84+
}
85+
description={item.comment}
86+
/>
87+
</List.Item>
88+
)}
89+
/>
90+
</Card>
91+
);
92+
};

‎frontend/javascripts/viewer/view/ai_jobs/alignment/ai_alignment_parameters.tsx‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import { SettingOutlined } from "@ant-design/icons";
2+
import { Card, Col, Form, Input, Row, Space } from "antd";
3+
import type { FormProps } from "antd";
4+
import type React from "react";
5+
import { ShouldUseManualMatchesFormItem } from "../../action-bar/ai_job_modals/components/should_use_trees_form_item";
6+
7+
import { useAlignmentJobContext } from "./ai_alignment_job_context";
8+
9+
export const AiAlignmentParameters: React.FC = () => {
10+
const { newDatasetName, setNewDatasetName, shouldUseManualMatches, setShouldUseManualMatches } =
11+
useAlignmentJobContext();
12+
13+
const handleValuesChange: FormProps["onValuesChange"] = (changedValues) => {
14+
if (Object.prototype.hasOwnProperty.call(changedValues, "newDatasetName")) {
15+
setNewDatasetName(changedValues.newDatasetName);
16+
}
17+
if (Object.prototype.hasOwnProperty.call(changedValues, "useAnnotation")) {
18+
setShouldUseManualMatches(changedValues.useAnnotation);
19+
}
20+
};
21+
22+
const formFields = [
23+
{ name: ["newDatasetName"], value: newDatasetName },
24+
{ name: ["useAnnotation"], value: shouldUseManualMatches },
25+
];
26+
27+
return (
28+
<Card
29+
title={
30+
<Space align="center">
31+
<SettingOutlined style={{ color: "#1890ff" }} />
32+
Alignment Parameters
33+
</Space>
34+
}
35+
>
36+
<Form layout="vertical" onValuesChange={handleValuesChange} fields={formFields}>
37+
<Row gutter={24}>
38+
<Col span={12}>
39+
<Form.Item
40+
name="newDatasetName"
41+
label="New Dataset Name"
42+
rules={[{ required: true, message: "Please provide a name for the new dataset" }]}
43+
>
44+
<Input />
45+
</Form.Item>
46+
</Col>
47+
<Col span={12}>
48+
<ShouldUseManualMatchesFormItem />
49+
</Col>
50+
</Row>
51+
</Form>
52+
</Card>
53+
);
54+
};

‎frontend/javascripts/viewer/view/ai_jobs/alignment/ai_image_alignment_job.tsx‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,21 @@
1-
import { Empty } from "antd";
1+
import { Flex } from "antd";
2+
import { AlignmentCreditInformation } from "../credit_information";
3+
import { AlignmentJobContextProvider } from "./ai_alignment_job_context";
4+
import { AiAlignmentModelSelector } from "./ai_alignment_model_selector";
5+
import { AiAlignmentParameters } from "./ai_alignment_parameters";
26

3-
export const AiImageAlignment = () => {
4-
return <Empty description="AI Image Alignment - Coming Soon" />;
7+
export const AiImageAlignmentJob = () => {
8+
return (
9+
<AlignmentJobContextProvider>
10+
<Flex gap={24}>
11+
<Flex flex="2" vertical gap={24}>
12+
<AiAlignmentModelSelector />
13+
<AiAlignmentParameters />
14+
</Flex>
15+
<Flex flex="1" vertical>
16+
<AlignmentCreditInformation />
17+
</Flex>
18+
</Flex>
19+
</AlignmentJobContextProvider>
20+
);
521
};

0 commit comments

Comments
 (0)