open-mmlab · Suntgr · Oct 30, 2024
diff --git a/packages/lb-annotation/src/constant/tool.ts b/packages/lb-annotation/src/constant/tool.ts
@@ -293,3 +293,12 @@ export enum EOperationMode {
   General = 1, // Common
   MultiMove = 2, // Experimental
 }
+
+/** rectToll name for LocalStorage */
+export const RECT_TOOL_MODE_NAME = 'rect_tool_mode';
+
+/** rectToll mode*/
+export enum ERectToolModeType {
+  ThreePoints = 'three_points',
+  TwoPoints = 'two_points',
+}
diff --git a/packages/lb-annotation/src/core/toolOperation/polygonOperation.ts b/packages/lb-annotation/src/core/toolOperation/polygonOperation.ts
@@ -11,7 +11,14 @@ import {
   TEXT_ATTRIBUTE_OFFSET,
 } from '../../constant/annotation';
 import EKeyCode from '../../constant/keyCode';
-import { edgeAdsorptionScope, ELineTypes, EPolygonPattern, EToolName } from '../../constant/tool';
+import {
+  edgeAdsorptionScope,
+  ELineTypes,
+  EPolygonPattern,
+  EToolName,
+  ERectToolModeType,
+  RECT_TOOL_MODE_NAME,
+} from '../../constant/tool';
 import locale from '../../locales';
 import { EMessage } from '../../locales/constants';
 import { IPolygonConfig, IPolygonData, IPolygonPoint } from '../../types/tool/polygon';
@@ -76,6 +83,8 @@ class PolygonOperation extends BasicToolOperation {
 
   public selection: Selection;
 
+  private rectToolMode?: ERectToolModeType;
+
   constructor(props: IPolygonOperationProps) {
     super(props);
     this.config = CommonToolUtils.jsonParser(props.config);
@@ -159,6 +168,18 @@ class PolygonOperation extends BasicToolOperation {
     return this.selectedPolygon?.textAttribute;
   }
 
+  public get isThreePointsMode() {
+    return this.pattern === EPolygonPattern.Rect && this.drawingPointList.length === 2;
+  }
+
+  public get isTwoPointsMode() {
+    return (
+      this.pattern === EPolygonPattern.Rect &&
+      this.rectToolMode === ERectToolModeType.TwoPoints &&
+      this.drawingPointList.length === 1
+    );
+  }
+
   // 是否直接执行操作
   public isAllowDouble = (e: MouseEvent) => {
     const { selectedID } = this;
@@ -325,6 +346,7 @@ class PolygonOperation extends BasicToolOperation {
       return;
     }
 
+    this.rectToolMode = localStorage.getItem(RECT_TOOL_MODE_NAME) as ERectToolModeType;
     this.deleteSelectedID();
     const coordinateZoom = this.getCoordinateUnderZoom(e);
     const coordinate = AxisUtils.changeDrawOutsideTarget(
@@ -358,9 +380,13 @@ class PolygonOperation extends BasicToolOperation {
       1 / this.zoom,
     );
 
-    if (this.pattern === EPolygonPattern.Rect && this.drawingPointList.length === 2) {
-      const rect = MathUtils.getRectangleByRightAngle(coordinateWithOrigin, this.drawingPointList);
-      this.drawingPointList = rect;
+    if (this.isThreePointsMode || this.isTwoPointsMode) {
+      if (this.isThreePointsMode) {
+        const rect = MathUtils.getRectangleByRightAngle(coordinateWithOrigin, this.drawingPointList);
+        this.drawingPointList = rect;
+      } else if (this.isTwoPointsMode) {
+        this.drawingPointList = this.createRectByTwoPointsMode(this.drawingPointList, coordinateWithOrigin);
+      }
 
       // 边缘判断 - 仅限支持图片下范围下
       if (this.config.drawOutsideTarget === false && this.imgInfo) {
@@ -1734,37 +1760,42 @@ class PolygonOperation extends BasicToolOperation {
       let drawingPointList = [...this.drawingPointList];
       let coordinate = AxisUtils.getOriginCoordinateWithOffsetCoordinate(this.coord, this.zoom, this.currentPos);
 
-      if (this.pattern === EPolygonPattern.Rect && drawingPointList.length === 2) {
+      if (this.isThreePointsMode) {
         // 矩形模式特殊绘制
         drawingPointList = MathUtils.getRectangleByRightAngle(coordinate, drawingPointList);
-      } else {
-        if (this.config?.edgeAdsorption && this.isAlt === false) {
-          const { dropFoot } = PolygonUtils.getClosestPoint(
-            coordinate,
-            this.polygonList,
-            this.config?.lineType,
-            edgeAdsorptionScope / this.zoom,
-          );
-          if (dropFoot) {
-            coordinate = dropFoot;
-          }
+      } else if (this.config?.edgeAdsorption && this.isAlt === false) {
+        const { dropFoot } = PolygonUtils.getClosestPoint(
+          coordinate,
+          this.polygonList,
+          this.config?.lineType,
+          edgeAdsorptionScope / this.zoom,
+        );
+        if (dropFoot) {
+          coordinate = dropFoot;
         }
         drawingPointList.push(coordinate);
+      } else if (this.isTwoPointsMode) {
+        drawingPointList = this.createRectByTwoPointsMode(drawingPointList, coordinate);
+      } else {
+        drawingPointList.push(coordinate);
+      }
+      const polygon = AxisUtils.changePointListByZoom(drawingPointList, this.zoom, this.currentPos);
+      DrawUtils.drawSelectedPolygonWithFillAndLine(this.canvas, polygon, {
+        fillColor: toolData.fill,
+        strokeColor: toolData.stroke,
+        pointColor: 'white',
+        thickness: 2,
+        lineCap: 'round',
+        isClose: false,
+        lineType: this.config.lineType,
+      });
+      if (this.isTwoPointsMode) {
+        DrawUtils.drawLine(this.canvas, polygon[0], polygon[1], {
+          color: 'white',
+          thickness: 3,
+          lineDash: [6],
+        });
       }
-
-      DrawUtils.drawSelectedPolygonWithFillAndLine(
-        this.canvas,
-        AxisUtils.changePointListByZoom(drawingPointList, this.zoom, this.currentPos),
-        {
-          fillColor: toolData.fill,
-          strokeColor: toolData.stroke,
-          pointColor: 'white',
-          thickness: 2,
-          lineCap: 'round',
-          isClose: false,
-          lineType: this.config.lineType,
-        },
-      );
     }
 
     // 5. 编辑中高亮的点
@@ -1805,6 +1836,30 @@ class PolygonOperation extends BasicToolOperation {
     }
   }
 
+  /**
+   * used in isTwoPointsMode；
+   * Effect：Determine the order of the four points, because the orientation is drawn along the line between the starting point and the first point；
+   * Orientation Logic: to the first point, adjacent to the two sides take one for orientation, such as side 1, Side 2, if side 1 clockwise rotation 90 degrees can get side 2, then take side 1 for orientation；
+   */
+  private createRectByTwoPointsMode(drawingPointList: IPolygonPoint[], coordinate: IPolygonPoint) {
+    const startPoint = drawingPointList[0];
+    const result = [...drawingPointList];
+    const fromLeftTopToRightBottom = coordinate.x > startPoint.x && coordinate.y > startPoint.y;
+    const fromRightBottomToLeftTop = coordinate.x < startPoint.x && coordinate.y < startPoint.y;
+    if (fromLeftTopToRightBottom || fromRightBottomToLeftTop) {
+      result.push({ x: startPoint.x, y: coordinate.y }, coordinate, {
+        x: coordinate.x,
+        y: startPoint.y,
+      });
+    } else {
+      result.push({ x: coordinate.x, y: startPoint.y }, coordinate, {
+        x: startPoint.x,
+        y: coordinate.y,
+      });
+    }
+    return result;
+  }
+
   public render() {
     if (!this.ctx) {
       return;

diff --git a/packages/lb-components/src/assets/annotation/rectTool/three_points_rect.svg b/packages/lb-components/src/assets/annotation/rectTool/three_points_rect.svg
diff --git a/packages/lb-components/src/assets/annotation/rectTool/two_points_rect.svg b/packages/lb-components/src/assets/annotation/rectTool/two_points_rect.svg
diff --git a/packages/lb-components/src/index.scss b/packages/lb-components/src/index.scss
@@ -309,6 +309,16 @@ $prefix: bee;
     margin-left: 16px;
   }
 
+  &__multiBox {
+    display: flex;
+    align-items: center;
+    margin: 0 5px;
+  }
+
+  &__dropdown .ant-dropdown-menu-item:not(.ant-dropdown-menu-item-selected) {
+    color: #999999;
+  }
+
   &__singleTool {
     max-width: 20px;
     max-height: 20px;

diff --git a/packages/lb-components/src/views/MainView/sidebar/ToolIcons.tsx b/packages/lb-components/src/views/MainView/sidebar/ToolIcons.tsx
@@ -13,19 +13,33 @@ import pointSvg from '@/assets/annotation/pointTool/icon_point.svg';
 import pointASvg from '@/assets/annotation/pointTool/icon_point_a.svg';
 import PolygonASvg from '@/assets/annotation/polygonTool/icon_polygon_a.svg';
 import PolygonSvg from '@/assets/annotation/polygonTool/icon_polygon.svg';
-import rectSvg from '@/assets/annotation/rectTool/icon_rect.svg';
-import rectASvg from '@/assets/annotation/rectTool/icon_rect_a.svg';
 import { cTool } from '@labelbee/lb-annotation';
 import classnames from 'classnames';
 import { useTranslation } from 'react-i18next';
+import type { MenuProps } from 'antd';
+import { Dropdown } from 'antd';
+import { DownOutlined } from '@ant-design/icons';
+import { ReactComponent as TwoPointsRectSvg } from '@/assets/annotation/rectTool/two_points_rect.svg';
+import { ReactComponent as ThreePointsRectSvg } from '@/assets/annotation/rectTool/three_points_rect.svg';
+import { useLocalStorageState } from 'ahooks';
 
-const { EPointCloudName, TOOL_NAME, TOOL_NAME_EN } = cTool;
+const { EPointCloudName, TOOL_NAME, TOOL_NAME_EN, ERectToolModeType, RECT_TOOL_MODE_NAME } = cTool;
 
 const toolList = [
   {
     toolName: EToolName.Rect,
-    commonSvg: rectSvg,
-    selectedSvg: rectASvg,
+    dropdownItems: [
+      {
+        key: ERectToolModeType.ThreePoints,
+        icon: <ThreePointsRectSvg />,
+        label: '三点画矩形',
+      },
+      {
+        key: ERectToolModeType.TwoPoints,
+        icon: <TwoPointsRectSvg />,
+        label: '两点画矩形',
+      },
+    ],
   },
   {
     toolName: EToolName.Polygon,
@@ -73,6 +87,19 @@ export const ToolIcons = ({
 
   const hasMultiTools = renderTools.length > 1;
 
+  const [rectToolMode, setRectToolMode] = useLocalStorageState(RECT_TOOL_MODE_NAME, {
+    defaultValue: ERectToolModeType.ThreePoints as string,
+    serializer: (v) => v ?? '',
+    deserializer: (v) => v,
+  });
+
+  const SelectedRectSvg =
+    rectToolMode === ERectToolModeType.TwoPoints ? TwoPointsRectSvg : ThreePointsRectSvg;
+
+  const toggleTwoOrThreePointRect: MenuProps['onClick'] = (e) => {
+    setRectToolMode(e.key);
+  };
+
   return (
     <div className={`${sidebarCls}__level`}>
       {renderTools.map((tool) => {
@@ -83,10 +110,30 @@ export const ToolIcons = ({
             key={tool.toolName}
             onClick={() => onChange?.(tool.toolName)}
           >
-            <img
-              className={`${sidebarCls}__singleTool`}
-              src={isSelected ? tool?.selectedSvg : tool?.commonSvg}
-            />
+            {tool.dropdownItems ? (
+              <Dropdown
+                overlayClassName={`${sidebarCls}__dropdown`}
+                menu={{
+                  items: tool.dropdownItems,
+                  selectable: true,
+                  defaultSelectedKeys: rectToolMode ? [rectToolMode] : [],
+                  onClick: toggleTwoOrThreePointRect,
+                }}
+              >
+                <div
+                  className={`${sidebarCls}__multiBox`}
+                  style={{ color: isSelected ? '#666fff' : '#999999' }}
+                >
+                  <SelectedRectSvg />
+                  <DownOutlined />
+                </div>
+              </Dropdown>
+            ) : (
+              <img
+                className={`${sidebarCls}__singleTool`}
+                src={isSelected ? tool?.selectedSvg : tool?.commonSvg}
+              />
+            )}
             <span
               className={classnames({
                 [`${sidebarCls}__toolOption__selected`]: isSelected,