From d2f83d0d54b256c22be0938a384aed582543e78c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roman=20R=C3=A4dle?= <raedle@meta.com>
Date: Tue, 24 Dec 2024 16:33:40 -0800
Subject: [PATCH] [sam2][demo][1/x] Remove adding object constraint after
 popagation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

PR #486 changed propagation to track each object independently. This allows to add objects after the initial propagation.

The frontend has a constraint that would prevent users from adding new objects after the initial "tracking". Now that the model supports tracking objects independently, we can remove this constraint from the UI.

Test Plan:

`cd demo/frontend`

`yarn lint`
```
(base) ➜ demo/frontend $ yarn lint                                                                                                                ⎇ remotes/origin/HEAD*
yarn run v1.18.0
$ eslint . --ext ts,tsx --report-unused-disable-directives --max-warnings 0
✨  Done in 16.98s.
```
---
 .../annotations/AddObjectButton.tsx           | 19 +++++++++++++++++--
 demo/frontend/src/demo/DemoConfig.tsx         |  3 +--
 demo/frontend/src/demo/atoms.ts               |  7 +------
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/demo/frontend/src/common/components/annotations/AddObjectButton.tsx b/demo/frontend/src/common/components/annotations/AddObjectButton.tsx
index ad52fa9..06a91d7 100644
--- a/demo/frontend/src/common/components/annotations/AddObjectButton.tsx
+++ b/demo/frontend/src/common/components/annotations/AddObjectButton.tsx
@@ -15,18 +15,33 @@
  */
 import useMessagesSnackbar from '@/common/components/snackbar/useDemoMessagesSnackbar';
 import useVideo from '@/common/components/video/editor/useVideo';
-import {activeTrackletObjectIdAtom, labelTypeAtom} from '@/demo/atoms';
+import {
+  activeTrackletObjectIdAtom,
+  isStreamingAtom,
+  labelTypeAtom,
+} from '@/demo/atoms';
 import {Add} from '@carbon/icons-react';
-import {useSetAtom} from 'jotai';
+import {useAtomValue, useSetAtom} from 'jotai';
 
 export default function AddObjectButton() {
   const video = useVideo();
   const setActiveTrackletId = useSetAtom(activeTrackletObjectIdAtom);
+  const isStreaming = useAtomValue(isStreamingAtom);
   const setLabelType = useSetAtom(labelTypeAtom);
   const {enqueueMessage} = useMessagesSnackbar();
 
   async function addObject() {
     enqueueMessage('addObjectClick');
+
+    // Make sure video is paused to allow users adding clicks. Either abort
+    // streaming if propagation is running or simply pause the video.
+    if (isStreaming) {
+      void video?.abortStreamMasks();
+    } else {
+      video?.pause();
+    }
+
+    // Create tracklet for video.
     const tracklet = await video?.createTracklet();
     if (tracklet != null) {
       setActiveTrackletId(tracklet.id);
diff --git a/demo/frontend/src/demo/DemoConfig.tsx b/demo/frontend/src/demo/DemoConfig.tsx
index 0816e5a..10a0398 100644
--- a/demo/frontend/src/demo/DemoConfig.tsx
+++ b/demo/frontend/src/demo/DemoConfig.tsx
@@ -24,8 +24,7 @@ export const DEMO_SHORT_NAME = 'SAM 2 Demo';
 export const RESEARCH_BY_META_AI = 'By Meta FAIR';
 export const DEMO_FRIENDLY_NAME = 'Segment Anything 2 Demo';
 export const VIDEO_WATERMARK_TEXT = `Modified with ${DEMO_FRIENDLY_NAME}`;
-export const PROJECT_GITHUB_URL =
-  'https://github.com/facebookresearch/sam2';
+export const PROJECT_GITHUB_URL = 'https://github.com/facebookresearch/sam2';
 export const AIDEMOS_URL = 'https://aidemos.meta.com';
 export const ABOUT_URL = 'https://ai.meta.com/sam2';
 export const EMAIL_ADDRESS = 'segment-anything@meta.com';
diff --git a/demo/frontend/src/demo/atoms.ts b/demo/frontend/src/demo/atoms.ts
index b7af65e..5f4dadc 100644
--- a/demo/frontend/src/demo/atoms.ts
+++ b/demo/frontend/src/demo/atoms.ts
@@ -123,14 +123,9 @@ export const pointsAtom = atom<SegmentationPoint[]>(get => {
 export const labelTypeAtom = atom<'positive' | 'negative'>('positive');
 
 export const isAddObjectEnabledAtom = atom<boolean>(get => {
-  const session = get(sessionAtom);
   const trackletsInitialized = get(areTrackletObjectsInitializedAtom);
   const isObjectLimitReached = get(isTrackletObjectLimitReachedAtom);
-  return (
-    session?.ranPropagation === false &&
-    trackletsInitialized &&
-    !isObjectLimitReached
-  );
+  return trackletsInitialized && !isObjectLimitReached;
 });
 
 export const codeEditorOpenedAtom = atom<boolean>(false);