Skip to content

Commit e29b9b4

Browse files
hotzenklotzhotzenklotz
committed
add first draft for new AI job runner
1 parent 6e19897 commit e29b9b4

File tree

11 files changed

+557
-4
lines changed

11 files changed

+557
-4
lines changed

‎frontend/javascripts/types/api_types.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -787,11 +787,11 @@ export enum APIJobType {
787787
FIND_LARGEST_SEGMENT_ID = "find_largest_segment_id",
788788
INFER_NUCLEI = "infer_nuclei",
789789
INFER_NEURONS = "infer_neurons",
790+
INFER_MITOCHONDRIA = "infer_mitochondria",
791+
INFER_INSTANCES = "infer_instances",
790792
MATERIALIZE_VOLUME_ANNOTATION = "materialize_volume_annotation",
791793
TRAIN_NEURON_MODEL = "train_neuron_model",
792794
TRAIN_INSTANCE_MODEL = "train_instance_model",
793-
INFER_MITOCHONDRIA = "infer_mitochondria",
794-
INFER_INSTANCES = "infer_instances",
795795
// Only used for backwards compatibility, e.g. to display results.
796796
DEPRECATED_INFER_WITH_MODEL = "infer_with_model",
797797
DEPRECATED_TRAIN_MODEL = "train_model",

‎frontend/javascripts/viewer/view/action_bar_view.tsx‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,11 @@ import {
4242
layoutEmitter,
4343
} from "viewer/view/layouting/layout_persistence";
4444
import type { StartAIJobModalState } from "./action-bar/ai_job_modals/constants";
45-
import { StartAIJobModal } from "./action-bar/ai_job_modals/start_ai_job_modal";
45+
// import { StartAIJobModal } from "./action-bar/ai_job_modals/start_ai_job_modal";
4646
import ToolkitView from "./action-bar/tools/toolkit_switcher_view";
4747
import ButtonComponent from "./components/button_component";
4848
import { NumberSliderSetting } from "./components/setting_input_views";
49+
import { AiJobsDrawer } from "./ai_jobs/ai_jobs_drawer";
4950

5051
const VersionRestoreWarning = (
5152
<Alert
@@ -376,7 +377,7 @@ class ActionBarView extends React.PureComponent<Props, State> {
376377
})
377378
}
378379
/>
379-
<StartAIJobModal aIJobModalState={this.props.aiJobModalState} />
380+
<AiJobsDrawer isOpen={this.props.aiJobModalState !== "invisible"} />
380381
</React.Fragment>
381382
);
382383
}

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { Drawer, Tabs } from "antd";
2+
import { AiImageSegmentation } from "./run_ai_model/ai_image_segmentation_job";
3+
import { AiModelTraining } from "./train_ai_model/ai_model_training_job";
4+
import { AiImageAlignment } from "./alignment/ai_image_alignment_job";
5+
6+
const { TabPane } = Tabs;
7+
8+
export const AiJobsDrawer = () => {
9+
return (
10+
<Drawer title="AI Jobs" placement="right" width={1200} open={true}>
11+
<Tabs defaultActiveKey="1">
12+
<TabPane tab="Image Segmentation" key="1">
13+
<AiImageSegmentation />
14+
</TabPane>
15+
<TabPane tab="Model Training" key="2">
16+
<AiModelTraining />
17+
</TabPane>
18+
<TabPane tab="Image Alignment" key="3">
19+
<AiImageAlignment />
20+
</TabPane>
21+
</Tabs>
22+
</Drawer>
23+
);
24+
};

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
import { Empty } from "antd";
2+
3+
export const AiImageAlignment = () => {
4+
return <Empty description="AI Image Alignment - Coming Soon" />;
5+
};

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import { useWkSelector } from "libs/react_hooks";
2+
import type React from "react";
3+
import { useMemo } from "react";
4+
import { getColorLayers } from "viewer/model/accessors/dataset_accessor";
5+
import { getUserBoundingBoxesFromState } from "viewer/model/accessors/tracing_accessor";
6+
import type { UserBoundingBox } from "viewer/store";
7+
import { BoundingBoxSelection } from "viewer/view/action-bar/ai_job_modals/components/bounding_box_selection";
8+
import { getBoundingBoxesForLayers } from "viewer/view/action-bar/ai_job_modals/utils";
9+
10+
interface BoundingBoxSelectorProps {
11+
value: UserBoundingBox | null;
12+
onChange: (value: UserBoundingBox | null) => void;
13+
}
14+
15+
export const BoundingBoxSelector: React.FC<BoundingBoxSelectorProps> = ({ value, onChange }) => {
16+
const userBoundingBoxes = useWkSelector(getUserBoundingBoxesFromState);
17+
const dataset = useWkSelector((state) => state.dataset);
18+
const colorLayers = getColorLayers(dataset);
19+
20+
const allBoundingBoxes = useMemo(() => {
21+
const defaultBBs = getBoundingBoxesForLayers(colorLayers);
22+
return defaultBBs.concat(userBoundingBoxes);
23+
}, [colorLayers, userBoundingBoxes]);
24+
25+
const handleChange = (id: number | null) => {
26+
const selected = allBoundingBoxes.find((bb) => bb.id === id) || null;
27+
onChange(selected);
28+
};
29+
30+
return (
31+
<BoundingBoxSelection
32+
userBoundingBoxes={allBoundingBoxes}
33+
setSelectedBoundingBoxId={handleChange}
34+
value={value?.id || null}
35+
showVolume
36+
/>
37+
);
38+
};

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

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import { type JobCreditCostInfo, getJobCreditCost } from "admin/rest_api";
2+
import { Button, Card, Col, Row, Spin, Typography } from "antd";
3+
import { formatCreditsString } from "libs/format_utils";
4+
import { useFetch } from "libs/react_helpers";
5+
import { useWkSelector } from "libs/react_hooks";
6+
import { computeArrayFromBoundingBox } from "libs/utils";
7+
import type React from "react";
8+
import { useRunAiModelJobContext } from "./run_ai_model/ai_image_segmentation_job_context";
9+
10+
const { Title, Text } = Typography;
11+
12+
export const CreditInformation: React.FC = () => {
13+
const { selectedModel, selectedJobType, selectedBoundingBox, handleStartAnalysis } =
14+
useRunAiModelJobContext();
15+
const organizationCredits = useWkSelector(
16+
(state) => state.activeOrganization?.creditBalance || "0",
17+
);
18+
19+
const jobCreditCostInfo = useFetch<JobCreditCostInfo | undefined>(
20+
async () =>
21+
selectedBoundingBox && selectedJobType
22+
? await getJobCreditCost(
23+
selectedJobType,
24+
computeArrayFromBoundingBox(selectedBoundingBox.boundingBox),
25+
)
26+
: undefined,
27+
undefined,
28+
[selectedBoundingBox, selectedJobType],
29+
);
30+
31+
const costInCredits = jobCreditCostInfo?.costInCredits;
32+
33+
return (
34+
<Card title="Credit Information">
35+
<Row justify="space-between" align="middle">
36+
<Col>
37+
<Text>Available Credits</Text>
38+
</Col>
39+
<Col>
40+
<Title level={2} style={{ margin: 0 }}>
41+
{formatCreditsString(organizationCredits)}
42+
</Title>
43+
</Col>
44+
</Row>
45+
<hr style={{ margin: "24px 0" }} />
46+
<Title level={5}>Cost Breakdown:</Title>
47+
<Row justify="space-between">
48+
<Col>
49+
<Text>Selected Model:</Text>
50+
</Col>
51+
<Col>
52+
<Text strong>{selectedModel ? selectedModel.name : "-"}</Text>
53+
</Col>
54+
</Row>
55+
<Row justify="space-between">
56+
<Col>
57+
<Text>Dataset Size (Est.):</Text>
58+
</Col>
59+
<Col>
60+
<Text strong>{selectedBoundingBox ? "Selected" : "-"}</Text>
61+
</Col>
62+
</Row>
63+
<hr style={{ margin: "24px 0" }} />
64+
<Row justify="space-between">
65+
<Col>
66+
<Text>Total Cost:</Text>
67+
</Col>
68+
<Col>
69+
{jobCreditCostInfo === undefined && selectedBoundingBox && selectedModel ? (
70+
<Spin size="small" />
71+
) : (
72+
<Text strong>{costInCredits ? formatCreditsString(costInCredits) : "FREE"}</Text>
73+
)}
74+
</Col>
75+
</Row>
76+
<Button
77+
type="primary"
78+
block
79+
size="large"
80+
style={{ marginTop: "24px" }}
81+
disabled={!selectedModel || !selectedBoundingBox}
82+
onClick={handleStartAnalysis}
83+
>
84+
Start Analysis
85+
</Button>
86+
</Card>
87+
);
88+
};

‎frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_parameters.tsx‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { Card, Col, Input, Row, Select, Typography } from "antd";
2+
import { useWkSelector } from "libs/react_hooks";
3+
import type React from "react";
4+
import { getColorLayers } from "viewer/model/accessors/dataset_accessor";
5+
import { BoundingBoxSelector } from "../bounding_box_selector";
6+
import { useRunAiModelJobContext } from "./ai_image_segmentation_job_context";
7+
8+
const { Text } = Typography;
9+
10+
export const AiAnalysisParameters: React.FC = () => {
11+
const {
12+
selectedBoundingBox,
13+
setSelectedBoundingBox,
14+
newDatasetName,
15+
setNewDatasetName,
16+
selectedLayerName,
17+
setSelectedLayerName,
18+
} = useRunAiModelJobContext();
19+
const dataset = useWkSelector((state) => state.dataset);
20+
const colorLayers = getColorLayers(dataset);
21+
22+
return (
23+
<Card title="Analysis Parameters">
24+
<Row gutter={24}>
25+
<Col span={12}>
26+
<div style={{ marginBottom: "16px" }}>
27+
<Text>New Dataset Name</Text>
28+
<Input value={newDatasetName} onChange={(e) => setNewDatasetName(e.target.value)} />
29+
</div>
30+
<div style={{ marginBottom: "16px" }}>
31+
<Text>Image Data Layer</Text>
32+
<Select
33+
style={{ width: "100%" }}
34+
value={selectedLayerName}
35+
onChange={setSelectedLayerName}
36+
options={colorLayers.map((l) => ({ value: l.name, label: l.name }))}
37+
/>
38+
</div>
39+
</Col>
40+
<Col span={12}>
41+
<div style={{ marginBottom: "16px" }}>
42+
<Text>Bounding Box</Text>
43+
<BoundingBoxSelector value={selectedBoundingBox} onChange={setSelectedBoundingBox} />
44+
</div>
45+
<div>
46+
<Text>Comments</Text>
47+
<Input.TextArea rows={4} />
48+
</div>
49+
</Col>
50+
</Row>
51+
</Card>
52+
);
53+
};

‎frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job.tsx‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import { Flex } from "antd";
2+
import { CreditInformation } from "../credit_information";
3+
import { AiAnalysisParameters } from "./ai_analysis_parameters";
4+
import { RunAiModelJobContextProvider } from "./ai_image_segmentation_job_context";
5+
import { AiModelSelector } from "./ai_model_selector";
6+
7+
export const AiImageSegmentation = () => {
8+
return (
9+
<RunAiModelJobContextProvider>
10+
<Flex gap={24}>
11+
<Flex flex="2" vertical gap={24}>
12+
<AiModelSelector />
13+
<AiAnalysisParameters />
14+
</Flex>
15+
<Flex flex="1" vertical>
16+
<CreditInformation />
17+
</Flex>
18+
</Flex>
19+
</RunAiModelJobContextProvider>
20+
);
21+
};

0 commit comments

Comments
 (0)