/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import assert from 'assert';
import { VSBuffer } from '../../../../../../base/common/buffer.js';
import { Mimes } from '../../../../../../base/common/mime.js';
import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
import { CellKind, IMainCellDto, IOutputDto, NotebookCellMetadata } from '../../../common/notebookCommon.js';
import { matchCellBasedOnSimilarties } from '../../../common/services/notebookCellMatching.js';
import { NotebookWorker } from '../../../common/services/notebookWebWorker.js';
import { URI } from '../../../../../../base/common/uri.js';
import { IDiffChange } from '../../../../../../base/common/diff/diff.js';


suite('NotebookDiff Diff Service', () => {
	ensureNoDisposablesAreLeakedInTestSuite();
	let worker: NotebookWorker;
	suiteSetup(() => {
		worker = new NotebookWorker();
	});
	suiteTeardown(() => {
		worker.dispose();
	});

	test('No changes', async () => {
		const { mapping, diff } = await mapCells([
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes.length, 0);
	});

	test('diff different source', async () => {
		const { mapping, diff } = await mapCells([
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange
		]);
	});

	test('diff different source (2 cells)', async () => {
		const { mapping, diff } = await mapCells([
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['z'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange
		]);
	});

	test('diff different source (first cell changed)', async () => {
		const { mapping, diff } = await mapCells([
			[['import sys'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['import pandas'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange
		]);
	});


	test('diff test small source', async () => {
		const { mapping, diff } = await mapCells([
			[['123456789'], 'javascript', CellKind.Code, [], {}]
		], [
			[['987654321'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange
		]);
	});

	test('diff test data single cell', async () => {
		const { mapping, diff } = await mapCells([
			[[
				'# This version has a bug\n',
				'def mult(a, b):\n',
				'    return a / b'
			], 'javascript', CellKind.Code, [], {}]
		], [
			[[
				'def mult(a, b):\n',
				'    \'This version is debugged.\'\n',
				'    return a * b'
			], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange
		]);
	});

	test('diff foo/foe (swapped cells)', async () => {
		const { mapping, diff } = await mapCells([
			[['def foe(x, y):\n', '    return x + y\n', 'foe(3, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([6])) }] }], { metadata: { collapsed: false }, executionOrder: 5 }],
			[['def foo(x, y):\n', '    return x * y\n', 'foo(1, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([2])) }] }], { metadata: { collapsed: false }, executionOrder: 6 }],
			[[''], 'javascript', CellKind.Code, [], {}]
		], [
			[['def foo(x, y):\n', '    return x * y\n', 'foo(1, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([2])) }] }], { metadata: { collapsed: false }, executionOrder: 6 }],
			[['def foe(x, y):\n', '    return x + y\n', 'foe(3, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([6])) }] }], { metadata: { collapsed: false }, executionOrder: 5 }],
			[[''], 'javascript', CellKind.Code, [], {}]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 2 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 1, originalLength: 1, modifiedStart: 2, modifiedLength: 0 } satisfies IDiffChange,
			// { originalStart: 1, originalLength: 1, modifiedStart: 2, modifiedLength: 0 } satisfies IDiffChange
		]);
	});

	test('diff foo/foe (swapped cells with changes in output)', async () => {
		const { mapping, diff } = await mapCells([
			[['def foe(x, y):\n', '    return x + y\n', 'foe(3, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([6])) }] }], { metadata: { collapsed: false }, executionOrder: 5 }],
			[['def foo(x, y):\n', '    return x * y\n', 'foo(1, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([2])) }] }], { metadata: { collapsed: false }, executionOrder: 6 }],
			[[''], 'javascript', CellKind.Code, [], {}]
		], [
			[['def foo(x, y):\n', '    return x * y\n', 'foo(1, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([6])) }] }], { metadata: { collapsed: false }, executionOrder: 5 }],
			[['def foe(x, y):\n', '    return x + y\n', 'foe(3, 2)'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([2])) }] }], { metadata: { collapsed: false }, executionOrder: 6 }],
			[[''], 'javascript', CellKind.Code, [], {}]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 2 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 0, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 1, originalLength: 1, modifiedStart: 2, modifiedLength: 0 } satisfies IDiffChange,
		]);
	});

	test('diff markdown', async () => {
		const { mapping, diff } = await mapCells([
			[['This is a test notebook with only markdown cells'], 'markdown', CellKind.Markup, [], {}],
			[['Lorem ipsum dolor sit amet'], 'markdown', CellKind.Markup, [], {}],
			[['In other news'], 'markdown', CellKind.Markup, [], {}],
		], [
			[['This is a test notebook with markdown cells only'], 'markdown', CellKind.Markup, [], {}],
			[['Lorem ipsum dolor sit amet'], 'markdown', CellKind.Markup, [], {}],
			[['In the news'], 'markdown', CellKind.Markup, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('diff insert 1 at the top', async () => {
		const { mapping, diff } = await mapCells([
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}]
		], [
			[['var h = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('diff insert 1 at the top with lots of cells', async () => {

		const { mapping, diff } = await mapCells([
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		], [
			[['var h = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 1 },
			{ modified: 3, original: 2 },
			{ modified: 4, original: 3 },
			{ modified: 5, original: 4 },
			{ modified: 6, original: 5 },
			{ modified: 7, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('diff insert 3 at the top with lots of cells', async () => {

		const { mapping, diff } = await mapCells([
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		], [
			[['var h = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var i = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var j = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 },
			{ modified: 1, original: -1 },
			{ modified: 2, original: -1 },
			{ modified: 3, original: 0 },
			{ modified: 4, original: 1 },
			{ modified: 5, original: 2 },
			{ modified: 6, original: 3 },
			{ modified: 7, original: 4 },
			{ modified: 8, original: 5 },
			{ modified: 9, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 3 } satisfies IDiffChange,
		]);
	});

	test('diff insert 1 in the middle', async () => {

		const { mapping, diff } = await mapCells([
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		], [
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var h = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: -1 },
			{ modified: 5, original: 4 },
			{ modified: 6, original: 5 },
			{ modified: 7, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 4, originalLength: 0, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('LCS (swap cells with outputs)', async () => {
		const { mapping, diff } = await mapCells([
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }]
		], [
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('LCS (outputs changed)', async () => {
		const { mapping, diff } = await mapCells([
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['y'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }]
		], [
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true } }],
			[['y'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Swap cells with output and with same code in each cell', async () => {
		const { mapping, diff } = await mapCells([
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
			[['x = 5'], 'javascript', CellKind.Code, [], {}],
			[['x'], 'javascript', CellKind.Code, [], {}],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], {}],
		], [
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['x = 5'], 'javascript', CellKind.Code, [], {}],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], {}],
			[['x'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 6, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Change outputs of cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['y'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
			[['x = 5'], 'javascript', CellKind.Code, [], {}],
			[['x'], 'javascript', CellKind.Code, [], {}],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], {}],
		], [
			[['# Description'], 'markdown', CellKind.Markup, [], { metadata: {} }],
			[['x = 3'], 'javascript', CellKind.Code, [], { metadata: { collapsed: true }, executionOrder: 1 }],
			[['x'], 'javascript', CellKind.Code, [], { metadata: { collapsed: false } }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 1 }],
			[['x = 5'], 'javascript', CellKind.Code, [], {}],
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], {}],
			[['y'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 6, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Insert a new cell', async () => {
		const { mapping, diff } = await mapCells([
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		], [
			[['var a = 1;'], 'javascript', CellKind.Code, [], {}],
			[['var b = 2;'], 'javascript', CellKind.Code, [], {}],
			[['var c = 3;'], 'javascript', CellKind.Code, [], {}],
			[['var d = 4;'], 'javascript', CellKind.Code, [], {}],
			[['var h = 8;'], 'javascript', CellKind.Code, [], {}],
			[['var e = 5;'], 'javascript', CellKind.Code, [], {}],
			[['var f = 6;'], 'javascript', CellKind.Code, [], {}],
			[['var g = 7;'], 'javascript', CellKind.Code, [], {}],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: -1 },//
			{ modified: 5, original: 4 },
			{ modified: 6, original: 5 },
			{ modified: 7, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 4, originalLength: 0, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('diff output', async () => {
		const { mapping, diff } = await mapCells([
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([4])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('diff output fast check', async () => {
		const { mapping, diff } = await mapCells([
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([4])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		], [
			[['x'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([3])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
			[['y'], 'javascript', CellKind.Code, [{ outputId: 'someOtherId', outputs: [{ mime: Mimes.text, data: VSBuffer.wrap(new Uint8Array([5])) }] }], { metadata: { collapsed: false }, executionOrder: 3 }],
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('No cells', async () => {
		const { mapping, diff } = await mapCells([
		], [
		]);

		assert.deepStrictEqual(mapping, [
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
		]);
	});
	test('No cells and add 1 cell', async () => {
		const { mapping, diff } = await mapCells([
		], [
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('One cells and delete 1 cell', async () => {
		const { mapping, diff } = await mapCells([
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
		]);

		assert.deepStrictEqual(mapping, [
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 0 } satisfies IDiffChange,
		]);
	});
	test('No changes with 1 cell', async () => {
		const { mapping, diff } = await mapCells([
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			// { originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 0 } satisfies IDiffChange,
		]);
	});
	test('One changes with 1 cell', async () => {
		const { mapping, diff } = await mapCells([
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		],
			[
				[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined]
			]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('No changes with 2 cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			// { originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('One changes with 2 cells, 1st cell changed', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('2 cells, 1st cell deleted', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 0 } satisfies IDiffChange,
		]);
	});
	test('2 cells, inserted a cell in the middle', async () => {
		const { mapping, diff } = await mapCells([
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined]
		], [
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['print("Bar Baz")'], 'python', CellKind.Code, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: -1 },
			{ modified: 2, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 0, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('One changes with 2 cells, 2nd cell changed', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		], [
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined]
		]);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Two cells, all cells changed', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined]
		]
			, [
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined]
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 2, modifiedStart: 0, modifiedLength: 2 } satisfies IDiffChange,
		]);
	});
	test('Insert a md cell and modify the next code cell (few cells)', async () => {
		const { mapping, diff } = await mapCells([
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Description'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello worldz")'], 'python', CellKind.Code, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 0, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Insert a md cell and modify the next code cell', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['print("bar baz1234")'], 'python', CellKind.Code, [], undefined]
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['# New Markdown cell'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo baz")'], 'python', CellKind.Code, [], undefined],
				[['print("bar baz1234")'], 'python', CellKind.Code, [], undefined]
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: -1 },
			{ modified: 4, original: 3 },
			{ modified: 5, original: 4 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 3, originalLength: 0, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Delete code cell and insert MD cell, both will not match', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['import sys'], 'python', CellKind.Code, [], undefined],
			[['sys.executable'], 'python', CellKind.Code, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['import sys'], 'python', CellKind.Code, [], undefined],
				[['# Header'], 'python', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: -1 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Delete MD cell and insert code cell, both will not match', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['import sys'], 'python', CellKind.Code, [], undefined],
			[['sys.executable'], 'python', CellKind.Code, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['import sys'], 'python', CellKind.Code, [], undefined],
				[['# Header'], 'python', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: -1 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Insert a md cell and modify the next 2 code cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined],
			[['print("bar baz")'], 'python', CellKind.Code, [], undefined]
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['# New Markdown cell'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Foo BaZ")'], 'python', CellKind.Code, [], undefined],
				[['print("bar baz Modified")'], 'python', CellKind.Code, [], undefined]
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: -1 },
			{ modified: 4, original: 3 },
			{ modified: 5, original: 4 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 3, originalLength: 0, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Insert a md cell, delete a cell and modify the next 2 code cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined],
			[['print("bar baz")'], 'python', CellKind.Code, [], undefined],
			[['print("Fox Trot")'], 'python', CellKind.Code, [], undefined]
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['# New Markdown cell'], 'markdown', CellKind.Markup, [], undefined],
				// [['print("Foo BaZ")'], 'python', CellKind.Code, [], undefined],
				[['print("bar baz Modified")'], 'python', CellKind.Code, [], undefined],
				[['print("Fox Trots")'], 'python', CellKind.Code, [], undefined]
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: -1 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Insert a md cell, delete a code cell and modify the next 2 code cells (deleted and inserted lineup)', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['import sys'], 'python', CellKind.Code, [], undefined],
			[['print("Foo Bar")'], 'python', CellKind.Code, [], undefined],
			[['print("bar baz")'], 'python', CellKind.Code, [], undefined],
			[['print("Fox Trot")'], 'python', CellKind.Code, [], undefined]
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['import sys'], 'python', CellKind.Code, [], undefined],
				[['sys.executable'], 'python', CellKind.Code, [], undefined],
				// [['print("Foo BaZ")'], 'python', CellKind.Code, [], undefined],
				[['print("bar baz Modified")'], 'python', CellKind.Code, [], undefined],
				[['print("Fox Trot")'], 'python', CellKind.Code, [], undefined]
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});

	test('Insert a few md cells and modify the next few code cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# New 1'], 'markdown', CellKind.Markup, [], undefined],
				[['# New 2'], 'markdown', CellKind.Markup, [], undefined],
				[['# Another MD cell (modified1)'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# Another MD cell (modified)'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: -1 },
			{ modified: 5, original: -1 },
			{ modified: 6, original: 4 },
			{ modified: 7, original: 5 },
			{ modified: 8, original: 6 },
			{ modified: 9, original: 7 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 4, originalLength: 0, modifiedStart: 4, modifiedLength: 2 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 6, originalLength: 1, modifiedStart: 8, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Insert & delete a few md cells and modify the next few code cells', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# New 1'], 'markdown', CellKind.Markup, [], undefined],
				[['# New 2'], 'markdown', CellKind.Markup, [], undefined],
				[['# Another MD cell (modified1)'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# Another MD cell (modified)'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: -1 },
			{ modified: 5, original: -1 },
			{ modified: 6, original: 4 },
			{ modified: 7, original: 5 },
			{ modified: 8, original: 6 },
			{ modified: 9, original: 7 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 4, originalLength: 0, modifiedStart: 4, modifiedLength: 2 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 6, originalLength: 1, modifiedStart: 8, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Modify, then insert ', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo BaZ'], 'markdown', CellKind.Markup, [], undefined],
				[['# Header'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
				[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
				[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: -1 },
			{ modified: 4, original: 3 },
			{ modified: 5, original: 4 },
			{ modified: 6, original: 5 },
			{ modified: 7, original: 6 },
			{ modified: 8, original: 7 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 0, modifiedStart: 3, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 6, originalLength: 1, modifiedStart: 8, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Modify, then delete ', async () => {
		const { mapping, diff } = await mapCells([
			[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
			[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
			[['# Foo Bar'], 'markdown', CellKind.Markup, [], undefined],
			[['print("foo bar")'], 'python', CellKind.Code, [], undefined],
			[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['import sys'], 'python', CellKind.Code, [], undefined],
			[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
			[['sys.executable'], 'python', CellKind.Code, [], undefined],
		]
			, [
				[['# Hello World'], 'markdown', CellKind.Markup, [], undefined],
				[['print("Hello world")'], 'python', CellKind.Code, [], undefined],
				[['# Foo BaZ'], 'markdown', CellKind.Markup, [], undefined],
				[['# Another MD cell'], 'markdown', CellKind.Markup, [], undefined],
				[['import sys'], 'python', CellKind.Code, [], undefined],
				[['# Yet another MD cell'], 'markdown', CellKind.Markup, [], undefined],
				[['sys.executable'], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 4 },
			{ modified: 4, original: 5 },
			{ modified: 5, original: 6 },
			{ modified: 6, original: 7 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 3, originalLength: 1, modifiedStart: 3, modifiedLength: 0 } satisfies IDiffChange,
			// { originalStart: 6, originalLength: 1, modifiedStart: 8, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Update code cell', async () => {
		const { mapping, diff } = await mapCells([
			[
				[
					"# This is a simple notebook\n",
					"\n",
					"There's nothing special here, I am just writing some text and plotting a function"
				], 'markdown', CellKind.Markup, [], undefined],
			[
				[
					"import numpy as np\n",
					"import matplotlib.pyplot as plt\n",
					"import pandas as pd\n",
					"\n",
					"%matplotlib inline"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"x = np.linspace(0, 4*np.pi,50)\n",
					"y = np.sin(x)\n",
					"\n",
					"plt.plot(x, y)"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"df = pd.DataFrame({f\"column_{c}\": np.random.normal(size=100) for c in range(100)})\n",
					"df"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"You could actually do some neat stuff with it, pandas like changing the colors (values above 1 should be green)"
				], 'python', CellKind.Markup, [], undefined],
			[
				[
					"df.style.applymap(lambda x: \"background-color: green\" if x>1 else \"background-color: white\")"
				], 'python', CellKind.Code, [], undefined],
			[
				[], 'python', CellKind.Code, [], undefined],
		]
			, [
				[
					[
						"# This is a simple notebook\n",
						"\n",
						"There's nothing special here, I am just writing some text and plotting a function"
					], 'markdown', CellKind.Markup, [], undefined],
				[
					[
						"import numpy as np\n",
						"import matplotlib.pyplot as plt\n",
						"import pandas as pd\n",
						"\n",
						"%matplotlib inline"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"x = np.linspace(0, 4*np.pi,50)\n",
						"y = 2 * np.sin(x)\n",
						"\n",
						"plt.plot(x, y)"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"df = pd.DataFrame({f\"column_{c}\": np.random.normal(size=100) for c in range(100)})\n",
						"df"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"You could actually do some neat stuff with it, pandas like changing the colors (values above 1 should be green)"
					], 'python', CellKind.Markup, [], undefined],
				[
					[
						"df.style.applymap(lambda x: \"background-color: green\" if x>1 else \"background-color: white\")"
					], 'python', CellKind.Code, [], undefined],
				[
					[], 'python', CellKind.Code, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
		]);

	});
	test('Update code cell 1', async () => {
		const { mapping, diff } = await mapCells([
			[
				[
					"# Hello World"
				], 'markdown', CellKind.Markup, [], undefined],
			[
				[
					"import os\n",
					"from pprint import pprint \n",
					"from sklearn.preprocessing import LabelEncoder\n",
					"import ast\n",
					"import pandas as pd\n",
					"from clipper import clipper\n",
					"import ffmpeg"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"data_dir = '../../../app/data/'\n",
					"output_dir = 'output/'\n",
					"video_dir = 'videos/'\n",
					"\n",
					"file_path = data_dir+output_dir+data_file\n",
					"video_path = data_dir+video_dir+video_file\n",
					"\n",
					"df = pd.read_csv(file_path)\n",
					"df.head()"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"# convert the string representation of results to a list\n",
					"df['list_categories'] = df['categories'].apply(ast.literal_eval)\n",
					"\n",
					"# get most likely label from each list\n",
					"df['label'] = df['list_categories'].apply(lambda x: x[0] if x else None)"
				], 'python', CellKind.Code, [], undefined],
			[
				[
					"# initialize the LabelEncoder\n",
					"label_encoder = LabelEncoder()\n",
					"\n",
					"# fit and transform the label to a numerical value\n",
					"numerical_data = label_encoder.fit_transform(df['label'])\n",
					"\n",
					"df['label_value'] = numerical_data"
				], 'python', CellKind.Markup, [], undefined],
		]
			, [
				[
					[
						"# Updated markdown cell"
					], 'markdown', CellKind.Markup, [], undefined],
				[
					[
						"from sklearn.preprocessing import LabelEncoder\n",
						"import ast\n",
						"import pandas as pd\n",
						"from clipper import clipper\n",
						"import ffmpeg"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"data_dir = '../../../app/data/'\n",
						"output_dir = 'output/'\n",
						"video_dir = 'videos/'\n",
						"\n",
						"file_path = data_dir+output_dir+data_file\n",
						"video_path = data_dir+video_dir+video_file\n",
						"\n",
						"df = pd.read_csv(file_path)\n",
						"df.head()"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"# convert the string representation of results to a list\n",
						"df['list_categories'] = df['categories'].apply(ast.literal_eval)\n",
						"\n",
						"# get most likely label from each list\n",
						"df['label'] = df['list_categories'].apply(lambda x: x[0] if x else None)"
					], 'python', CellKind.Code, [], undefined],
				[
					[
						"# initialize the LabelEncoder\n",
						"label_encoder = LabelEncoder()\n",
						"\n",
						"# fit and transform the label to a numerical value\n",
						"numerical_data = label_encoder.fit_transform(df['label'])\n",
						"\n",
						"df['label_value'] = numerical_data"
					], 'python', CellKind.Markup, [], undefined],
			]
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 1, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
		]);
	});
	test('Detect modification and insertion of cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"Make sure your [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) is working properly.\n",
						"To avoid errors that may arise from the version of the dependency package, it is recommended to use the **JupyterLab** instead of the Jupyter notebook to display image results.\n",
						"```\n",
						"- pip install --upgrade pip && pip install -r requirements.txt\n",
						"- jupyter labextension install --no-build @jupyter-widgets/jupyterlab-manager\n",
						"- jupyter labextension install --no-build jupyter-datawidgets/extension\n",
						"- jupyter labextension install jupyter-threejs\n",
						"- jupyter labextension list\n",
						"```\n",
						"\n",
						"You should see:\n",
						"```\n",
						"JupyterLab v...\n",
						"  ...\n",
						"    jupyterlab-datawidgets v... enabled OK\n",
						"    @jupyter-widgets/jupyterlab-manager v... enabled OK\n",
						"    jupyter-threejs v... enabled OK\n",
						"```\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/3D-pose-estimation-webcam/3D-pose-estimation.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using a Jupyter Notebook release. Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino-dev>=2024.0.0\" \"opencv-python\" \"torch\" \"onnx<1.16.2\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"with open(\"notebook_utils.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/engine3js.py\",\n",
						")\n",
						"with open(\"engine3js.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"import notebook_utils as utils\n",
						"import engine3js as engine"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We use `omz_downloader`, which is a command line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# directory where model will be downloaded\n",
						"base_model_dir = \"model\"\n",
						"\n",
						"# model name as named in Open Model Zoo\n",
						"model_name = \"human-pose-estimation-3d-0001\"\n",
						"# selected precision (FP32, FP16)\n",
						"precision = \"FP32\"\n",
						"\n",
						"BASE_MODEL_NAME = f\"{base_model_dir}/public/{model_name}/{model_name}\"\n",
						"model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
						"onnx_path = Path(BASE_MODEL_NAME).with_suffix(\".onnx\")\n",
						"\n",
						"ir_model_path = Path(f\"model/public/{model_name}/{precision}/{model_name}.xml\")\n",
						"model_weights_path = Path(f\"model/public/{model_name}/{precision}/{model_name}.bin\")\n",
						"\n",
						"if not model_path.exists():\n",
						"    download_command = f\"omz_downloader \" f\"--name {model_name} \" f\"--output_dir {base_model_dir}\"\n",
						"    ! $download_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The selected model comes from the public directory, which means it must be converted into OpenVINO Intermediate Representation (OpenVINO IR). We use `omz_converter` to convert the ONNX format model to the OpenVINO IR format."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"if not onnx_path.exists():\n",
						"    convert_command = (\n",
						"        f\"omz_converter \" f\"--name {model_name} \" f\"--precisions {precision} \" f\"--download_dir {base_model_dir} \" f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $convert_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"device = utils.device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(model=ir_model_path, weights=model_weights_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)\n",
						"infer_request = compiled_model.create_infer_request()\n",
						"input_tensor_name = model.inputs[0].get_any_name()\n",
						"\n",
						"# get input and output names of nodes\n",
						"input_layer = compiled_model.input(0)\n",
						"output_layers = list(compiled_model.outputs)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The input for the model is data from the input image and the outputs are heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"input_layer.any_name, [o.any_name for o in output_layers]"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[None,]\n",
						"    infer_request.infer({input_tensor_name: img})\n",
						"    # A set of three inference results is obtained\n",
						"    results = {name: infer_request.get_tensor(name).data[:] for name in {\"features\", \"heatmaps\", \"pafs\"}}\n",
						"    # Get the results\n",
						"    results = (results[\"features\"][0], results[\"heatmaps\"][0], results[\"pafs\"][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1],\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15],\n",
						"        [15, 16],  # nose - l_eye - l_ear\n",
						"        [1, 17],\n",
						"        [17, 18],  # nose - r_eye - r_ear\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],  # neck - nose\n",
						"        [1, 16],\n",
						"        [16, 18],  # nose - l_eye - l_ear\n",
						"        [1, 15],\n",
						"        [15, 17],  # nose - r_eye - r_ear\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(source, flip=flip, fps=30, skip_first_frames=skip_frames)\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(grid=True, axis=True, view_width=windows_width, view_height=windows_height)\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(format=\"jpg\", height=windows_height, width=windows_width)\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2]))\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = engine.parse_poses(inference_result, 1, stride, focal_length, True)\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://storage.openvinotoolkit.org/data/test_data/videos/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
			, [
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"Make sure your [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) is working properly.\n",
						"To avoid errors that may arise from the version of the dependency package, it is recommended to use the **JupyterLab** instead of the Jupyter notebook to display image results.\n",
						"```\n",
						"- pip install --upgrade pip && pip install -r requirements.txt\n",
						"- jupyter labextension install --no-build @jupyter-widgets/jupyterlab-manager\n",
						"- jupyter labextension install --no-build jupyter-datawidgets/extension\n",
						"- jupyter labextension install jupyter-threejs\n",
						"- jupyter labextension list\n",
						"```\n",
						"\n",
						"You should see:\n",
						"```\n",
						"JupyterLab v...\n",
						"  ...\n",
						"    jupyterlab-datawidgets v... enabled OK\n",
						"    @jupyter-widgets/jupyterlab-manager v... enabled OK\n",
						"    jupyter-threejs v... enabled OK\n",
						"```\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/3D-pose-estimation-webcam/3D-pose-estimation.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using a Jupyter Notebook release. Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino>=2024.4.0\" \"opencv-python\" \"torch\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"with open(\"notebook_utils.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/engine3js.py\",\n",
						")\n",
						"with open(\"engine3js.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"import notebook_utils as utils\n",
						"import engine3js as engine"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"from notebook_utils import download_file\n",
						"import tarfile\n",
						"\n",
						"\n",
						"# directory where model will be downloaded\n",
						"base_model_dir = Path(\"model\")\n",
						"\n",
						"download_file(\"https://storage.openvinotoolkit.org/repositories/open_model_zoo/public/2022.1/human-pose-estimation-3d-0001/human-pose-estimation-3d.tar.gz\", directory=base_model_dir)\n",
						"\n",
						"ckpt_file =  base_model_dir / \"human-pose-estimation-3d-0001.pth\"\n",
						"\n",
						"if not ckpt_file.exists():\n",
						"    with tarfile.open(base_model_dir / \"human-pose-estimation-3d.tar.gz\") as f:\n",
						"        f.extractall(base_model_dir)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import torch\n",
						"import openvino as ov\n",
						"\n",
						"ov_model_path = Path(base_model_dir) / \"human-pose-estimation-3d-0001.xml\"\n",
						"\n",
						"if not ov_model_path.exists():\n",
						"    from model.model import PoseEstimationWithMobileNet\n",
						"\n",
						"    pose_estimation_model = PoseEstimationWithMobileNet(is_convertible_by_mo=True)\n",
						"    pose_estimation_model.loas_state_dict(torch.load(ckpt_file, map_location=\"cpu\"))\n",
						"    pose_estimation_model.eval()\n",
						"\n",
						"    with torch.no_grad():\n",
						"        ov_model = ov.convert_model(pose_estimation_model, example_input=torch.zeros([1, 3, 256, 448]), input=[1, 3, 256, 448])\n",
						"        ov.save_model(ov_model, ov_model_path)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"device = utils.device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(ov_model_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    mean_value = 128.0\n",
						"    scale_value = 255.0\n",
						"\n",
						"    img = (img - mean_value) / scale_value\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[None,]\n",
						"    result = compiled_model(img)\n",
						"    # Get the results\n",
						"    results = (result[0][0], results[1][0], results[2][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1],\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15],\n",
						"        [15, 16],  # nose - l_eye - l_ear\n",
						"        [1, 17],\n",
						"        [17, 18],  # nose - r_eye - r_ear\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],  # neck - nose\n",
						"        [1, 16],\n",
						"        [16, 18],  # nose - l_eye - l_ear\n",
						"        [1, 15],\n",
						"        [15, 17],  # nose - r_eye - r_ear\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(source, flip=flip, fps=30, skip_first_frames=skip_frames)\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(grid=True, axis=True, view_width=windows_width, view_height=windows_height)\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(format=\"jpg\", height=windows_height, width=windows_width)\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2]))\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = engine.parse_poses(inference_result, 1, stride, focal_length, True)\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://storage.openvinotoolkit.org/data/test_data/videos/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 15 },
			{ modified: 14, original: 16 },
			{ modified: 15, original: 17 },
			{ modified: 16, original: 18 },
			{ modified: 17, original: 19 },
			{ modified: 18, original: 20 },
			{ modified: 19, original: 21 },
			{ modified: 20, original: 22 },
			{ modified: 21, original: 23 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 2, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
			{
				modifiedLength: 1,
				modifiedStart: 6,
				originalLength: 1,
				originalStart: 6
			},
			{
				modifiedLength: 1,
				modifiedStart: 7,
				originalLength: 1,
				originalStart: 7
			},
			{
				modifiedLength: 1,
				modifiedStart: 8,
				originalLength: 1,
				originalStart: 8
			},
			{
				modifiedLength: 1,
				modifiedStart: 12,
				originalLength: 1,
				originalStart: 12
			},
			{
				modifiedLength: 0,
				modifiedStart: 13,
				originalLength: 2,
				originalStart: 13
			},
			{
				modifiedLength: 1,
				modifiedStart: 14,
				originalLength: 1,
				originalStart: 16
			}
		]);
	});
	test('Detect modification in 2 cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using the latest Jupyter Notebook release (2.4.1). Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino-dev>=2024.0.0\" \"opencv-python\" \"torch\" \"onnx\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import sys\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"import notebook_utils as utils\n",
						"\n",
						"sys.path.append(\"./engine\")\n",
						"import engine.engine3js as engine\n",
						"from engine.parse_poses import parse_poses"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We use `omz_downloader`, which is a command line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# directory where model will be downloaded\n",
						"base_model_dir = \"model\"\n",
						"\n",
						"# model name as named in Open Model Zoo\n",
						"model_name = \"human-pose-estimation-3d-0001\"\n",
						"# selected precision (FP32, FP16)\n",
						"precision = \"FP32\"\n",
						"\n",
						"BASE_MODEL_NAME = f\"{base_model_dir}/public/{model_name}/{model_name}\"\n",
						"model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
						"onnx_path = Path(BASE_MODEL_NAME).with_suffix(\".onnx\")\n",
						"\n",
						"ir_model_path = f\"model/public/{model_name}/{precision}/{model_name}.xml\"\n",
						"model_weights_path = f\"model/public/{model_name}/{precision}/{model_name}.bin\"\n",
						"\n",
						"if not model_path.exists():\n",
						"    download_command = f\"omz_downloader \" f\"--name {model_name} \" f\"--output_dir {base_model_dir}\"\n",
						"    ! $download_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The selected model comes from the public directory, which means it must be converted into OpenVINO Intermediate Representation (OpenVINO IR). We use `omz_converter` to convert the ONNX format model to the OpenVINO IR format."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"if not onnx_path.exists():\n",
						"    convert_command = (\n",
						"        f\"omz_converter \" f\"--name {model_name} \" f\"--precisions {precision} \" f\"--download_dir {base_model_dir} \" f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $convert_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()\n",
						"\n",
						"device = widgets.Dropdown(\n",
						"    options=core.available_devices + [\"AUTO\"],\n",
						"    value=\"AUTO\",\n",
						"    description=\"Device:\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(model=ir_model_path, weights=model_weights_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)\n",
						"infer_request = compiled_model.create_infer_request()\n",
						"input_tensor_name = model.inputs[0].get_any_name()\n",
						"\n",
						"# get input and output names of nodes\n",
						"input_layer = compiled_model.input(0)\n",
						"output_layers = list(compiled_model.outputs)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The input for the model is data from the input image and the outputs are heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"input_layer.any_name, [o.any_name for o in output_layers]"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[None,]\n",
						"    infer_request.infer({input_tensor_name: img})\n",
						"    # A set of three inference results is obtained\n",
						"    results = {name: infer_request.get_tensor(name).data[:] for name in {\"features\", \"heatmaps\", \"pafs\"}}\n",
						"    # Get the results\n",
						"    results = (results[\"features\"][0], results[\"heatmaps\"][0], results[\"pafs\"][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1],\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15],\n",
						"        [15, 16],  # nose - l_eye - l_ear\n",
						"        [1, 17],\n",
						"        [17, 18],  # nose - r_eye - r_ear\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],  # neck - nose\n",
						"        [1, 16],\n",
						"        [16, 18],  # nose - l_eye - l_ear\n",
						"        [1, 15],\n",
						"        [15, 17],  # nose - r_eye - r_ear\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(source, flip=flip, fps=30, skip_first_frames=skip_frames)\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(grid=True, axis=True, view_width=windows_width, view_height=windows_height)\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(format=\"jpg\", height=windows_height, width=windows_width)\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2]))\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = parse_poses(inference_result, 1, stride, focal_length, True)\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://github.com/intel-iot-devkit/sample-videos/raw/master/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
			, [
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using a Jupyter Notebook release. Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino-dev>=2024.0.0\" \"opencv-python\" \"torch\" \"onnx\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"with open(\"notebook_utils.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/engine3js.py\",\n",
						")\n",
						"with open(\"engine3js.py\", \"w\") as f:\n",
						"    f.write(r.text)\n",
						"\n",
						"import notebook_utils as utils\n",
						"import engine3js as engine"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We use `omz_downloader`, which is a command line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# directory where model will be downloaded\n",
						"base_model_dir = \"model\"\n",
						"\n",
						"# model name as named in Open Model Zoo\n",
						"model_name = \"human-pose-estimation-3d-0001\"\n",
						"# selected precision (FP32, FP16)\n",
						"precision = \"FP32\"\n",
						"\n",
						"BASE_MODEL_NAME = f\"{base_model_dir}/public/{model_name}/{model_name}\"\n",
						"model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
						"onnx_path = Path(BASE_MODEL_NAME).with_suffix(\".onnx\")\n",
						"\n",
						"ir_model_path = f\"model/public/{model_name}/{precision}/{model_name}.xml\"\n",
						"model_weights_path = f\"model/public/{model_name}/{precision}/{model_name}.bin\"\n",
						"\n",
						"if not model_path.exists():\n",
						"    download_command = f\"omz_downloader \" f\"--name {model_name} \" f\"--output_dir {base_model_dir}\"\n",
						"    ! $download_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The selected model comes from the public directory, which means it must be converted into OpenVINO Intermediate Representation (OpenVINO IR). We use `omz_converter` to convert the ONNX format model to the OpenVINO IR format."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"if not onnx_path.exists():\n",
						"    convert_command = (\n",
						"        f\"omz_converter \" f\"--name {model_name} \" f\"--precisions {precision} \" f\"--download_dir {base_model_dir} \" f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $convert_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()\n",
						"\n",
						"device = widgets.Dropdown(\n",
						"    options=core.available_devices + [\"AUTO\"],\n",
						"    value=\"AUTO\",\n",
						"    description=\"Device:\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(model=ir_model_path, weights=model_weights_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)\n",
						"infer_request = compiled_model.create_infer_request()\n",
						"input_tensor_name = model.inputs[0].get_any_name()\n",
						"\n",
						"# get input and output names of nodes\n",
						"input_layer = compiled_model.input(0)\n",
						"output_layers = list(compiled_model.outputs)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The input for the model is data from the input image and the outputs are heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"input_layer.any_name, [o.any_name for o in output_layers]"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[None,]\n",
						"    infer_request.infer({input_tensor_name: img})\n",
						"    # A set of three inference results is obtained\n",
						"    results = {name: infer_request.get_tensor(name).data[:] for name in {\"features\", \"heatmaps\", \"pafs\"}}\n",
						"    # Get the results\n",
						"    results = (results[\"features\"][0], results[\"heatmaps\"][0], results[\"pafs\"][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1],\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15],\n",
						"        [15, 16],  # nose - l_eye - l_ear\n",
						"        [1, 17],\n",
						"        [17, 18],  # nose - r_eye - r_ear\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],  # neck - nose\n",
						"        [1, 16],\n",
						"        [16, 18],  # nose - l_eye - l_ear\n",
						"        [1, 15],\n",
						"        [15, 17],  # nose - r_eye - r_ear\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(source, flip=flip, fps=30, skip_first_frames=skip_frames)\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(grid=True, axis=True, view_width=windows_width, view_height=windows_height)\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(format=\"jpg\", height=windows_height, width=windows_width)\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2]))\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = engine.parse_poses(inference_result, 1, stride, focal_length, True)\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://github.com/intel-iot-devkit/sample-videos/raw/master/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 1, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 4, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 20, originalLength: 1, modifiedStart: 20, modifiedLength: 1 } satisfies IDiffChange,
		]);

	});
	test('Detect modification in multiple cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using the latest Jupyter Notebook release (2.4.1). Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino-dev>=2024.0.0\" \"opencv-python\" \"torch\" \"onnx\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import sys\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"r = requests.get(\n",
						"    url='https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py',\n",
						")\n",
						"open('notebook_utils.py', 'w').write(r.text)\n",
						"import notebook_utils as utils\n",
						"\n",
						"sys.path.append(\"./engine\")\n",
						"import engine.engine3js as engine\n",
						"from engine.parse_poses import parse_poses"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We use `omz_downloader`, which is a command line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# directory where model will be downloaded\n",
						"base_model_dir = \"model\"\n",
						"\n",
						"# model name as named in Open Model Zoo\n",
						"model_name = \"human-pose-estimation-3d-0001\"\n",
						"# selected precision (FP32, FP16)\n",
						"precision = \"FP32\"\n",
						"\n",
						"BASE_MODEL_NAME = f\"{base_model_dir}/public/{model_name}/{model_name}\"\n",
						"model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
						"onnx_path = Path(BASE_MODEL_NAME).with_suffix(\".onnx\")\n",
						"\n",
						"ir_model_path = f\"model/public/{model_name}/{precision}/{model_name}.xml\"\n",
						"model_weights_path = f\"model/public/{model_name}/{precision}/{model_name}.bin\"\n",
						"\n",
						"if not model_path.exists():\n",
						"    download_command = (\n",
						"        f\"omz_downloader \" f\"--name {model_name} \" f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $download_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The selected model comes from the public directory, which means it must be converted into OpenVINO Intermediate Representation (OpenVINO IR). We use `omz_converter` to convert the ONNX format model to the OpenVINO IR format."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"if not onnx_path.exists():\n",
						"    convert_command = (\n",
						"        f\"omz_converter \"\n",
						"        f\"--name {model_name} \"\n",
						"        f\"--precisions {precision} \"\n",
						"        f\"--download_dir {base_model_dir} \"\n",
						"        f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $convert_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()\n",
						"\n",
						"device = widgets.Dropdown(\n",
						"    options=core.available_devices + [\"AUTO\"],\n",
						"    value='AUTO',\n",
						"    description='Device:',\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(model=ir_model_path, weights=model_weights_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)\n",
						"infer_request = compiled_model.create_infer_request()\n",
						"input_tensor_name = model.inputs[0].get_any_name()\n",
						"\n",
						"# get input and output names of nodes\n",
						"input_layer = compiled_model.input(0)\n",
						"output_layers = list(compiled_model.outputs)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The input for the model is data from the input image and the outputs are heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"input_layer.any_name, [o.any_name for o in output_layers]"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[\n",
						"        None,\n",
						"    ]\n",
						"    infer_request.infer({input_tensor_name: img})\n",
						"    # A set of three inference results is obtained\n",
						"    results = {\n",
						"        name: infer_request.get_tensor(name).data[:]\n",
						"        for name in {\"features\", \"heatmaps\", \"pafs\"}\n",
						"    }\n",
						"    # Get the results\n",
						"    results = (results[\"features\"][0], results[\"heatmaps\"][0], results[\"pafs\"][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1], \n",
						"        [0, 9], [9, 10], [10, 11],    # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3], [3, 4], [4, 5],       # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15], [15, 16],            # nose - l_eye - l_ear\n",
						"        [1, 17], [17, 18],            # nose - r_eye - r_ear\n",
						"        [0, 6], [6, 7], [7, 8],       # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12], [12, 13], [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],                       # neck - nose\n",
						"        [1, 16], [16, 18],            # nose - l_eye - l_ear\n",
						"        [1, 15], [15, 17],            # nose - r_eye - r_ear\n",
						"        [0, 3], [3, 4], [4, 5],       # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9], [9, 10], [10, 11],    # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6], [6, 7], [7, 8],       # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12], [12, 13], [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")  \n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(source, flip=flip, fps=30, skip_first_frames=skip_frames)\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(grid=True, axis=True, view_width=windows_width, view_height=windows_height)\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(\n",
						"                format=\"jpg\", height=windows_height, width=windows_width\n",
						"            )\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2]))\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = parse_poses(inference_result, 1, stride, focal_length, True)\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://github.com/intel-iot-devkit/sample-videos/raw/master/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
			, [
				{
					"cell_type": "markdown",
					"source": [
						"# Live 3D Human Pose Estimation with OpenVINO\n",
						"\n",
						"This notebook demonstrates live 3D Human Pose Estimation with OpenVINO via a webcam. We utilize the model [human-pose-estimation-3d-0001](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the end of this notebook, you will see live inference results from your webcam (if available). Alternatively, you can also upload a video file to test out the algorithms.\n",
						"**Make sure you have properly installed the [Jupyter extension](https://github.com/jupyter-widgets/pythreejs#jupyterlab) and been using JupyterLab to run the demo as suggested in the `README.md`**\n",
						"\n",
						"> **NOTE**: _To use a webcam, you must run this Jupyter notebook on a computer with a webcam. If you run on a remote server, the webcam will not work. However, you can still do inference on a video file in the final step. This demo utilizes the Python interface in `Three.js` integrated with WebGL to process data from the model inference. These results are processed and displayed in the notebook._\n",
						"\n",
						"_To ensure that the results are displayed correctly, run the code in a recommended browser on one of the following operating systems:_\n",
						"_Ubuntu, Windows: Chrome_\n",
						"_macOS: Safari_\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Imports](#Imports)\n",
						"- [The model](#The-model)\n",
						"    - [Download the model](#Download-the-model)\n",
						"    - [Convert Model to OpenVINO IR format](#Convert-Model-to-OpenVINO-IR-format)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Load the model](#Load-the-model)\n",
						"- [Processing](#Processing)\n",
						"    - [Model Inference](#Model-Inference)\n",
						"    - [Draw 2D Pose Overlays](#Draw-2D-Pose-Overlays)\n",
						"    - [Main Processing Function](#Main-Processing-Function)\n",
						"- [Run](#Run)\n",
						"\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"**The `pythreejs` extension may not display properly when using the latest Jupyter Notebook release (2.4.1). Therefore, it is recommended to use Jupyter Lab instead.**"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install pythreejs \"openvino-dev>=2024.0.0\" \"opencv-python\" \"torch\" \"onnx\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Imports\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import collections\n",
						"import sys\n",
						"import time\n",
						"from pathlib import Path\n",
						"\n",
						"import cv2\n",
						"import ipywidgets as widgets\n",
						"import numpy as np\n",
						"from IPython.display import clear_output, display\n",
						"import openvino as ov\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"import notebook_utils as utils\n",
						"\n",
						"sys.path.append(\"./engine\")\n",
						"import engine.engine3js as engine\n",
						"from engine.parse_poses import parse_poses"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## The model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Download the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We use `omz_downloader`, which is a command line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# directory where model will be downloaded\n",
						"base_model_dir = \"model\"\n",
						"\n",
						"# model name as named in Open Model Zoo\n",
						"model_name = \"human-pose-estimation-3d-0001\"\n",
						"# selected precision (FP32, FP16)\n",
						"precision = \"FP32\"\n",
						"\n",
						"BASE_MODEL_NAME = f\"{base_model_dir}/public/{model_name}/{model_name}\"\n",
						"model_path = Path(BASE_MODEL_NAME).with_suffix(\".pth\")\n",
						"onnx_path = Path(BASE_MODEL_NAME).with_suffix(\".onnx\")\n",
						"\n",
						"ir_model_path = f\"model/public/{model_name}/{precision}/{model_name}.xml\"\n",
						"model_weights_path = f\"model/public/{model_name}/{precision}/{model_name}.bin\"\n",
						"\n",
						"if not model_path.exists():\n",
						"    download_command = (\n",
						"        f\"omz_downloader \" f\"--name {model_name} \" f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $download_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Convert Model to OpenVINO IR format\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The selected model comes from the public directory, which means it must be converted into OpenVINO Intermediate Representation (OpenVINO IR). We use `omz_converter` to convert the ONNX format model to the OpenVINO IR format."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"if not onnx_path.exists():\n",
						"    convert_command = (\n",
						"        f\"omz_converter \"\n",
						"        f\"--name {model_name} \"\n",
						"        f\"--precisions {precision} \"\n",
						"        f\"--download_dir {base_model_dir} \"\n",
						"        f\"--output_dir {base_model_dir}\"\n",
						"    )\n",
						"    ! $convert_command"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()\n",
						"\n",
						"device = widgets.Dropdown(\n",
						"    options=core.available_devices + [\"AUTO\"],\n",
						"    value=\"AUTO\",\n",
						"    description=\"Device:\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Converted models are located in a fixed structure, which indicates vendor, model name and precision.\n",
						"\n",
						"First, initialize the inference engine, OpenVINO Runtime. Then, read the network architecture and model weights from the `.bin` and `.xml` files to compile for the desired device. An inference request is then created to infer the compiled model."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# initialize inference engine\n",
						"core = ov.Core()\n",
						"# read the network and corresponding weights from file\n",
						"model = core.read_model(model=ir_model_path, weights=model_weights_path)\n",
						"# load the model on the specified device\n",
						"compiled_model = core.compile_model(model=model, device_name=device.value)\n",
						"infer_request = compiled_model.create_infer_request()\n",
						"input_tensor_name = model.inputs[0].get_any_name()\n",
						"\n",
						"# get input and output names of nodes\n",
						"input_layer = compiled_model.input(0)\n",
						"output_layers = list(compiled_model.outputs)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The input for the model is data from the input image and the outputs are heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"input_layer.any_name, [o.any_name for o in output_layers]"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Processing\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"### Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Frames captured from video files or the live webcam are used as the input for the 3D model. This is how you obtain the output heat maps, PAF (part affinity fields) and features."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def model_infer(scaled_img, stride):\n",
						"    \"\"\"\n",
						"    Run model inference on the input image\n",
						"\n",
						"    Parameters:\n",
						"        scaled_img: resized image according to the input size of the model\n",
						"        stride: int, the stride of the window\n",
						"    \"\"\"\n",
						"\n",
						"    # Remove excess space from the picture\n",
						"    img = scaled_img[\n",
						"        0 : scaled_img.shape[0] - (scaled_img.shape[0] % stride),\n",
						"        0 : scaled_img.shape[1] - (scaled_img.shape[1] % stride),\n",
						"    ]\n",
						"\n",
						"    img = np.transpose(img, (2, 0, 1))[None,]\n",
						"    infer_request.infer({input_tensor_name: img})\n",
						"    # A set of three inference results is obtained\n",
						"    results = {\n",
						"        name: infer_request.get_tensor(name).data[:]\n",
						"        for name in {\"features\", \"heatmaps\", \"pafs\"}\n",
						"    }\n",
						"    # Get the results\n",
						"    results = (results[\"features\"][0], results[\"heatmaps\"][0], results[\"pafs\"][0])\n",
						"\n",
						"    return results"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Draw 2D Pose Overlays\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We need to define some connections between the joints in advance, so that we can draw the structure of the human body in the resulting image after obtaining the inference results.\n",
						"Joints are drawn as circles and limbs are drawn as lines. The code is based on the [3D Human Pose Estimation Demo](https://github.com/openvinotoolkit/open_model_zoo/tree/master/demos/human_pose_estimation_3d_demo/python) from Open Model Zoo."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# 3D edge index array\n",
						"body_edges = np.array(\n",
						"    [\n",
						"        [0, 1],\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [1, 15],\n",
						"        [15, 16],  # nose - l_eye - l_ear\n",
						"        [1, 17],\n",
						"        [17, 18],  # nose - r_eye - r_ear\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"body_edges_2d = np.array(\n",
						"    [\n",
						"        [0, 1],  # neck - nose\n",
						"        [1, 16],\n",
						"        [16, 18],  # nose - l_eye - l_ear\n",
						"        [1, 15],\n",
						"        [15, 17],  # nose - r_eye - r_ear\n",
						"        [0, 3],\n",
						"        [3, 4],\n",
						"        [4, 5],  # neck - l_shoulder - l_elbow - l_wrist\n",
						"        [0, 9],\n",
						"        [9, 10],\n",
						"        [10, 11],  # neck - r_shoulder - r_elbow - r_wrist\n",
						"        [0, 6],\n",
						"        [6, 7],\n",
						"        [7, 8],  # neck - l_hip - l_knee - l_ankle\n",
						"        [0, 12],\n",
						"        [12, 13],\n",
						"        [13, 14],  # neck - r_hip - r_knee - r_ankle\n",
						"    ]\n",
						")\n",
						"\n",
						"\n",
						"def draw_poses(frame, poses_2d, scaled_img, use_popup):\n",
						"    \"\"\"\n",
						"    Draw 2D pose overlays on the image to visualize estimated poses.\n",
						"    Joints are drawn as circles and limbs are drawn as lines.\n",
						"\n",
						"    :param frame: the input image\n",
						"    :param poses_2d: array of human joint pairs\n",
						"    \"\"\"\n",
						"    for pose in poses_2d:\n",
						"        pose = np.array(pose[0:-1]).reshape((-1, 3)).transpose()\n",
						"        was_found = pose[2] > 0\n",
						"\n",
						"        pose[0], pose[1] = (\n",
						"            pose[0] * frame.shape[1] / scaled_img.shape[1],\n",
						"            pose[1] * frame.shape[0] / scaled_img.shape[0],\n",
						"        )\n",
						"\n",
						"        # Draw joints.\n",
						"        for edge in body_edges_2d:\n",
						"            if was_found[edge[0]] and was_found[edge[1]]:\n",
						"                cv2.line(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, edge[0]].astype(np.int32)),\n",
						"                    tuple(pose[0:2, edge[1]].astype(np.int32)),\n",
						"                    (255, 255, 0),\n",
						"                    4,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"        # Draw limbs.\n",
						"        for kpt_id in range(pose.shape[1]):\n",
						"            if pose[2, kpt_id] != -1:\n",
						"                cv2.circle(\n",
						"                    frame,\n",
						"                    tuple(pose[0:2, kpt_id].astype(np.int32)),\n",
						"                    3,\n",
						"                    (0, 255, 255),\n",
						"                    -1,\n",
						"                    cv2.LINE_AA,\n",
						"                )\n",
						"\n",
						"    return frame"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Main Processing Function\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run 3D pose estimation on the specified source. It could be either a webcam feed or a video file."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"def run_pose_estimation(source=0, flip=False, use_popup=False, skip_frames=0):\n",
						"    \"\"\"\n",
						"    2D image as input, using OpenVINO as inference backend,\n",
						"    get joints 3D coordinates, and draw 3D human skeleton in the scene\n",
						"\n",
						"    :param source:      The webcam number to feed the video stream with primary webcam set to \"0\", or the video path.\n",
						"    :param flip:        To be used by VideoPlayer function for flipping capture image.\n",
						"    :param use_popup:   False for showing encoded frames over this notebook, True for creating a popup window.\n",
						"    :param skip_frames: Number of frames to skip at the beginning of the video.\n",
						"    \"\"\"\n",
						"\n",
						"    focal_length = -1  # default\n",
						"    stride = 8\n",
						"    player = None\n",
						"    skeleton_set = None\n",
						"\n",
						"    try:\n",
						"        # create video player to play with target fps  video_path\n",
						"        # get the frame from camera\n",
						"        # You can skip first N frames to fast forward video. change 'skip_first_frames'\n",
						"        player = utils.VideoPlayer(\n",
						"            source, flip=flip, fps=30, skip_first_frames=skip_frames\n",
						"        )\n",
						"        # start capturing\n",
						"        player.start()\n",
						"\n",
						"        input_image = player.next()\n",
						"        # set the window size\n",
						"        resize_scale = 450 / input_image.shape[1]\n",
						"        windows_width = int(input_image.shape[1] * resize_scale)\n",
						"        windows_height = int(input_image.shape[0] * resize_scale)\n",
						"\n",
						"        # use visualization library\n",
						"        engine3D = engine.Engine3js(\n",
						"            grid=True, axis=True, view_width=windows_width, view_height=windows_height\n",
						"        )\n",
						"\n",
						"        if use_popup:\n",
						"            # display the 3D human pose in this notebook, and origin frame in popup window\n",
						"            display(engine3D.renderer)\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(title, cv2.WINDOW_KEEPRATIO | cv2.WINDOW_AUTOSIZE)\n",
						"        else:\n",
						"            # set the 2D image box, show both human pose and image in the notebook\n",
						"            imgbox = widgets.Image(\n",
						"                format=\"jpg\", height=windows_height, width=windows_width\n",
						"            )\n",
						"            display(widgets.HBox([engine3D.renderer, imgbox]))\n",
						"\n",
						"        skeleton = engine.Skeleton(body_edges=body_edges)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"\n",
						"        while True:\n",
						"            # grab the frame\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"\n",
						"            # resize image and change dims to fit neural network input\n",
						"            # (see https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/human-pose-estimation-3d-0001)\n",
						"            scaled_img = cv2.resize(\n",
						"                frame, dsize=(model.inputs[0].shape[3], model.inputs[0].shape[2])\n",
						"            )\n",
						"\n",
						"            if focal_length < 0:  # Focal length is unknown\n",
						"                focal_length = np.float32(0.8 * scaled_img.shape[1])\n",
						"\n",
						"            # inference start\n",
						"            start_time = time.time()\n",
						"            # get results\n",
						"            inference_result = model_infer(scaled_img, stride)\n",
						"\n",
						"            # inference stop\n",
						"            stop_time = time.time()\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Process the point to point coordinates of the data\n",
						"            poses_3d, poses_2d = parse_poses(\n",
						"                inference_result, 1, stride, focal_length, True\n",
						"            )\n",
						"\n",
						"            # use processing times from last 200 frames\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"\n",
						"            if len(poses_3d) > 0:\n",
						"                # From here, you can rotate the 3D point positions using the function \"draw_poses\",\n",
						"                # or you can directly make the correct mapping below to properly display the object image on the screen\n",
						"                poses_3d_copy = poses_3d.copy()\n",
						"                x = poses_3d_copy[:, 0::4]\n",
						"                y = poses_3d_copy[:, 1::4]\n",
						"                z = poses_3d_copy[:, 2::4]\n",
						"                poses_3d[:, 0::4], poses_3d[:, 1::4], poses_3d[:, 2::4] = (\n",
						"                    -z + np.ones(poses_3d[:, 2::4].shape) * 200,\n",
						"                    -y + np.ones(poses_3d[:, 2::4].shape) * 100,\n",
						"                    -x,\n",
						"                )\n",
						"\n",
						"                poses_3d = poses_3d.reshape(poses_3d.shape[0], 19, -1)[:, :, 0:3]\n",
						"                people = skeleton(poses_3d=poses_3d)\n",
						"\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"                engine3D.scene_add(people)\n",
						"                skeleton_set = people\n",
						"\n",
						"                # draw 2D\n",
						"                frame = draw_poses(frame, poses_2d, scaled_img, use_popup)\n",
						"\n",
						"            else:\n",
						"                try:\n",
						"                    engine3D.scene_remove(skeleton_set)\n",
						"                    skeleton_set = None\n",
						"                except Exception:\n",
						"                    pass\n",
						"\n",
						"            cv2.putText(\n",
						"                frame,\n",
						"                f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                (10, 30),\n",
						"                cv2.FONT_HERSHEY_COMPLEX,\n",
						"                0.7,\n",
						"                (0, 0, 255),\n",
						"                1,\n",
						"                cv2.LINE_AA,\n",
						"            )\n",
						"\n",
						"            if use_popup:\n",
						"                cv2.imshow(title, frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27, use ESC to exit\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # encode numpy array to jpg\n",
						"                imgbox.value = cv2.imencode(\n",
						"                    \".jpg\",\n",
						"                    frame,\n",
						"                    params=[cv2.IMWRITE_JPEG_QUALITY, 90],\n",
						"                )[1].tobytes()\n",
						"\n",
						"            engine3D.renderer.render(engine3D.scene, engine3D.cam)\n",
						"\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        clear_output()\n",
						"        if player is not None:\n",
						"            # stop capturing\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()\n",
						"        if skeleton_set:\n",
						"            engine3D.scene_remove(skeleton_set)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Run\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Run, using a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						"> **NOTE**:\n",
						">\n",
						"> *1. To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a server (e.g. Binder), the webcam will not work.*\n",
						">\n",
						"> *2. Popup mode may not work if you run this notebook on a remote computer (e.g. Binder).*\n",
						"\n",
						"If you do not have a webcam, you can still run this demo with a video file. Any [format supported by OpenCV](https://docs.opencv.org/4.5.1/dd/d43/tutorial_py_video_display.html) will work."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Using the following method, you can click and move your mouse over the picture on the left to interact."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"tags": []
					},
					"source": [
						"USE_WEBCAM = False\n",
						"\n",
						"cam_id = 0\n",
						"video_path = \"https://github.com/intel-iot-devkit/sample-videos/raw/master/face-demographics-walking.mp4\"\n",
						"\n",
						"source = cam_id if USE_WEBCAM else video_path\n",
						"\n",
						"run_pose_estimation(source=source, flip=isinstance(source, int), use_popup=False)"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 4, originalLength: 1, modifiedStart: 4, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 10, originalLength: 1, modifiedStart: 10, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 16, originalLength: 1, modifiedStart: 16, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 18, originalLength: 1, modifiedStart: 18, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 20, originalLength: 1, modifiedStart: 20, modifiedLength: 1 } satisfies IDiffChange,
		]);

	});
	test('Modification with Insertion, detected as deletion of a cell', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Video generation with ZeroScope and OpenVINO\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Install and import required packages](#Install-and-import-required-packages)\n",
						"- [Load the model](#Load-the-model)\n",
						"- [Convert the model](#Convert-the-model)\n",
						"    - [Define the conversion function](#Define-the-conversion-function)\n",
						"    - [UNet](#UNet)\n",
						"    - [VAE](#VAE)\n",
						"    - [Text encoder](#Text-encoder)\n",
						"- [Build a pipeline](#Build-a-pipeline)\n",
						"- [Inference with OpenVINO](#Inference-with-OpenVINO)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Define a prompt](#Define-a-prompt)\n",
						"    - [Video generation](#Video-generation)\n",
						"- [Interactive demo](#Interactive-demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/zeroscope-text2video/zeroscope-text2video.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The ZeroScope model is a free and open-source text-to-video model that can generate realistic and engaging videos from text descriptions. It is based on the [Modelscope](https://modelscope.cn/models/damo/text-to-video-synthesis/summary) model, but it has been improved to produce higher-quality videos with a 16:9 aspect ratio and no Shutterstock watermark. The ZeroScope model is available in two versions: ZeroScope_v2 576w, which is optimized for rapid content creation at a resolution of 576x320 pixels, and ZeroScope_v2 XL, which upscales videos to a high-definition resolution of 1024x576.\n",
						"\n",
						"The ZeroScope model is trained on a dataset of over 9,000 videos and 29,000 tagged frames. It uses a diffusion model to generate videos, which means that it starts with a random noise image and gradually adds detail to it until it matches the text description. The ZeroScope model is still under development, but it has already been used to create some impressive videos. For example, it has been used to create videos of people dancing, playing sports, and even driving cars.\n",
						"\n",
						"The ZeroScope model is a powerful tool that can be used to create various videos, from simple animations to complex scenes. It is still under development, but it has the potential to revolutionize the way we create and consume video content.\n",
						"\n",
						"Both versions of the ZeroScope model are available on Hugging Face:\n",
						" - [ZeroScope_v2 576w](https://huggingface.co/cerspense/zeroscope_v2_576w)\n",
						" - [ZeroScope_v2 XL](https://huggingface.co/cerspense/zeroscope_v2_XL)\n",
						"\n",
						"We will use the first one."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"<div class=\"alert alert-block alert-warning\">\n",
						"    This tutorial requires at least 24GB of free memory to generate a video with a frame size of 432x240 and 16 frames. Increasing either of these values will require more memory and take more time.\n",
						"</div>"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Install and import required packages\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"To work with text-to-video synthesis model, we will use Hugging Face's [Diffusers](https://github.com/huggingface/diffusers) library. It provides already pretrained model from `cerspense`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu  \"diffusers>=0.18.0\" \"torch>=2.1\" transformers \"openvino>=2023.1.0\" numpy \"gradio>=4.19\""
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import gc\n",
						"from typing import Optional, Union, List, Callable\n",
						"import base64\n",
						"import tempfile\n",
						"import warnings\n",
						"\n",
						"import diffusers\n",
						"import transformers\n",
						"import numpy as np\n",
						"import IPython\n",
						"import torch\n",
						"import PIL\n",
						"import gradio as gr\n",
						"\n",
						"import openvino as ov"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Original 576x320 inference requires a lot of RAM (>100GB), so let's run our example on a smaller frame size, keeping the same aspect ratio. Try reducing values below to reduce the memory consumption."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"WIDTH = 432  # must be divisible by 8\n",
						"HEIGHT = 240  # must be divisible by 8\n",
						"NUM_FRAMES = 16"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The model is loaded from HuggingFace using `.from_pretrained` method of `diffusers.DiffusionPipeline`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"pipe = diffusers.DiffusionPipeline.from_pretrained(\"cerspense/zeroscope_v2_576w\")"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet = pipe.unet\n",
						"unet.eval()\n",
						"vae = pipe.vae\n",
						"vae.eval()\n",
						"text_encoder = pipe.text_encoder\n",
						"text_encoder.eval()\n",
						"tokenizer = pipe.tokenizer\n",
						"scheduler = pipe.scheduler\n",
						"vae_scale_factor = pipe.vae_scale_factor\n",
						"unet_in_channels = pipe.unet.config.in_channels\n",
						"sample_width = WIDTH // vae_scale_factor\n",
						"sample_height = HEIGHT // vae_scale_factor\n",
						"del pipe\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Convert the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The architecture for generating videos from text comprises three distinct sub-networks: one for extracting text features, another for translating text features into the video latent space using a diffusion model, and a final one for mapping the video latent space to the visual space. The collective parameters of the entire model amount to approximately 1.7 billion. It's capable of processing English input. The diffusion model is built upon the Unet3D model and achieves video generation by iteratively denoising a starting point of pure Gaussian noise video."
					]
				},
				{
					"cell_type": "markdown",
					"source": []
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define the conversion function\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Model components are PyTorch modules, that can be converted with `ov.convert_model` function directly. We also use `ov.save_model` function to serialize the result of conversion."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"\n",
						"def convert(model: torch.nn.Module, xml_path: str, **convert_kwargs) -> Path:\n",
						"    xml_path = Path(xml_path)\n",
						"    if not xml_path.exists():\n",
						"        xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
						"        with torch.no_grad():\n",
						"            converted_model = ov.convert_model(model, **convert_kwargs)\n",
						"        ov.save_model(converted_model, xml_path)\n",
						"        del converted_model\n",
						"        gc.collect()\n",
						"        torch._C._jit_clear_class_registry()\n",
						"        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
						"        torch.jit._state._clear_class_state()\n",
						"    return xml_path"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### UNet\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text-to-video generation pipeline main component is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet_xml_path = convert(\n",
						"    unet,\n",
						"    \"models/unet.xml\",\n",
						"    example_input={\n",
						"        \"sample\": torch.randn(2, 4, 2, int(sample_height // 2), int(sample_width // 2)),\n",
						"        \"timestep\": torch.tensor(1),\n",
						"        \"encoder_hidden_states\": torch.randn(2, 77, 1024),\n",
						"    },\n",
						"    input=[\n",
						"        (\"sample\", (2, 4, NUM_FRAMES, sample_height, sample_width)),\n",
						"        (\"timestep\", ()),\n",
						"        (\"encoder_hidden_states\", (2, 77, 1024)),\n",
						"    ],\n",
						")\n",
						"del unet\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### VAE\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Variational autoencoder (VAE) uses UNet output to decode latents to visual representations. Our VAE model has KL loss for encoding images into latents and decoding latent representations into images. For inference, we need only decoder part."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"class VaeDecoderWrapper(torch.nn.Module):\n",
						"    def __init__(self, vae):\n",
						"        super().__init__()\n",
						"        self.vae = vae\n",
						"\n",
						"    def forward(self, z: torch.FloatTensor):\n",
						"        return self.vae.decode(z)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"vae_decoder_xml_path = convert(\n",
						"    VaeDecoderWrapper(vae),\n",
						"    \"models/vae.xml\",\n",
						"    example_input=torch.randn(2, 4, 32, 32),\n",
						"    input=((NUM_FRAMES, 4, sample_height, sample_width)),\n",
						")\n",
						"del vae\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Text encoder\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text encoder is used to encode the input prompt to tensor. Default tensor length is 77."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"text_encoder_xml = convert(\n",
						"    text_encoder,\n",
						"    \"models/text_encoder.xml\",\n",
						"    example_input=torch.ones(1, 77, dtype=torch.int64),\n",
						"    input=((1, 77), ov.Type.i64),\n",
						")\n",
						"del text_encoder\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Build a pipeline\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> List[np.ndarray]:\n",
						"    # This code is copied from https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78\n",
						"    # reshape to ncfhw\n",
						"    mean = torch.tensor(mean, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    std = torch.tensor(std, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    # unnormalize back to [0,1]\n",
						"    video = video.mul_(std).add_(mean)\n",
						"    video.clamp_(0, 1)\n",
						"    # prepare the final outputs\n",
						"    i, c, f, h, w = video.shape\n",
						"    images = video.permute(2, 3, 0, 4, 1).reshape(f, h, i * w, c)  # 1st (frames, h, batch_size, w, c) 2nd (frames, h, batch_size * w, c)\n",
						"    images = images.unbind(dim=0)  # prepare a list of indvidual (consecutive frames)\n",
						"    images = [(image.cpu().numpy() * 255).astype(\"uint8\") for image in images]  # f h w c\n",
						"    return images"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"try:\n",
						"    from diffusers.utils import randn_tensor\n",
						"except ImportError:\n",
						"    from diffusers.utils.torch_utils import randn_tensor\n",
						"\n",
						"\n",
						"class OVTextToVideoSDPipeline(diffusers.DiffusionPipeline):\n",
						"    def __init__(\n",
						"        self,\n",
						"        vae_decoder: ov.CompiledModel,\n",
						"        text_encoder: ov.CompiledModel,\n",
						"        tokenizer: transformers.CLIPTokenizer,\n",
						"        unet: ov.CompiledModel,\n",
						"        scheduler: diffusers.schedulers.DDIMScheduler,\n",
						"    ):\n",
						"        super().__init__()\n",
						"\n",
						"        self.vae_decoder = vae_decoder\n",
						"        self.text_encoder = text_encoder\n",
						"        self.tokenizer = tokenizer\n",
						"        self.unet = unet\n",
						"        self.scheduler = scheduler\n",
						"        self.vae_scale_factor = vae_scale_factor\n",
						"        self.unet_in_channels = unet_in_channels\n",
						"        self.width = WIDTH\n",
						"        self.height = HEIGHT\n",
						"        self.num_frames = NUM_FRAMES\n",
						"\n",
						"    def __call__(\n",
						"        self,\n",
						"        prompt: Union[str, List[str]] = None,\n",
						"        num_inference_steps: int = 50,\n",
						"        guidance_scale: float = 9.0,\n",
						"        negative_prompt: Optional[Union[str, List[str]]] = None,\n",
						"        eta: float = 0.0,\n",
						"        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n",
						"        latents: Optional[torch.FloatTensor] = None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        output_type: Optional[str] = \"np\",\n",
						"        return_dict: bool = True,\n",
						"        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n",
						"        callback_steps: int = 1,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Function invoked when calling the pipeline for generation.\n",
						"\n",
						"        Args:\n",
						"            prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.\n",
						"                instead.\n",
						"            num_inference_steps (`int`, *optional*, defaults to 50):\n",
						"                The number of denoising steps. More denoising steps usually lead to a higher quality videos at the\n",
						"                expense of slower inference.\n",
						"            guidance_scale (`float`, *optional*, defaults to 7.5):\n",
						"                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).\n",
						"                `guidance_scale` is defined as `w` of equation 2. of [Imagen\n",
						"                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n",
						"                1`. Higher guidance scale encourages to generate videos that are closely linked to the text `prompt`,\n",
						"                usually at the expense of lower video quality.\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the video generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            eta (`float`, *optional*, defaults to 0.0):\n",
						"                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n",
						"                [`schedulers.DDIMScheduler`], will be ignored for others.\n",
						"            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n",
						"                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n",
						"                to make generation deterministic.\n",
						"            latents (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video\n",
						"                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n",
						"                tensor will ge generated by sampling using the supplied random `generator`. Latents should be of shape\n",
						"                `(batch_size, num_channel, num_frames, height, width)`.\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"            output_type (`str`, *optional*, defaults to `\"np\"`):\n",
						"                The output format of the generate video. Choose between `torch.FloatTensor` or `np.array`.\n",
						"            return_dict (`bool`, *optional*, defaults to `True`):\n",
						"                Whether or not to return a [`~pipelines.stable_diffusion.TextToVideoSDPipelineOutput`] instead of a\n",
						"                plain tuple.\n",
						"            callback (`Callable`, *optional*):\n",
						"                A function that will be called every `callback_steps` steps during inference. The function will be\n",
						"                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n",
						"            callback_steps (`int`, *optional*, defaults to 1):\n",
						"                The frequency at which the `callback` function will be called. If not specified, the callback will be\n",
						"                called at every step.\n",
						"\n",
						"        Returns:\n",
						"            `List[np.ndarray]`: generated video frames\n",
						"        \"\"\"\n",
						"\n",
						"        num_images_per_prompt = 1\n",
						"\n",
						"        # 1. Check inputs. Raise error if not correct\n",
						"        self.check_inputs(\n",
						"            prompt,\n",
						"            callback_steps,\n",
						"            negative_prompt,\n",
						"            prompt_embeds,\n",
						"            negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 2. Define call parameters\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n",
						"        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n",
						"        # corresponds to doing no classifier free guidance.\n",
						"        do_classifier_free_guidance = guidance_scale > 1.0\n",
						"\n",
						"        # 3. Encode input prompt\n",
						"        prompt_embeds = self._encode_prompt(\n",
						"            prompt,\n",
						"            num_images_per_prompt,\n",
						"            do_classifier_free_guidance,\n",
						"            negative_prompt,\n",
						"            prompt_embeds=prompt_embeds,\n",
						"            negative_prompt_embeds=negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 4. Prepare timesteps\n",
						"        self.scheduler.set_timesteps(num_inference_steps)\n",
						"        timesteps = self.scheduler.timesteps\n",
						"\n",
						"        # 5. Prepare latent variables\n",
						"        num_channels_latents = self.unet_in_channels\n",
						"        latents = self.prepare_latents(\n",
						"            batch_size * num_images_per_prompt,\n",
						"            num_channels_latents,\n",
						"            prompt_embeds.dtype,\n",
						"            generator,\n",
						"            latents,\n",
						"        )\n",
						"\n",
						"        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline\n",
						"        extra_step_kwargs = {\"generator\": generator, \"eta\": eta}\n",
						"\n",
						"        # 7. Denoising loop\n",
						"        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n",
						"        with self.progress_bar(total=num_inference_steps) as progress_bar:\n",
						"            for i, t in enumerate(timesteps):\n",
						"                # expand the latents if we are doing classifier free guidance\n",
						"                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n",
						"                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
						"\n",
						"                # predict the noise residual\n",
						"                noise_pred = self.unet(\n",
						"                    {\n",
						"                        \"sample\": latent_model_input,\n",
						"                        \"timestep\": t,\n",
						"                        \"encoder_hidden_states\": prompt_embeds,\n",
						"                    }\n",
						"                )[0]\n",
						"                noise_pred = torch.tensor(noise_pred)\n",
						"\n",
						"                # perform guidance\n",
						"                if do_classifier_free_guidance:\n",
						"                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
						"                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
						"\n",
						"                # reshape latents\n",
						"                bsz, channel, frames, width, height = latents.shape\n",
						"                latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"                noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"\n",
						"                # compute the previous noisy sample x_t -> x_t-1\n",
						"                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample\n",
						"\n",
						"                # reshape latents back\n",
						"                latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4)\n",
						"\n",
						"                # call the callback, if provided\n",
						"                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n",
						"                    progress_bar.update()\n",
						"                    if callback is not None and i % callback_steps == 0:\n",
						"                        callback(i, t, latents)\n",
						"\n",
						"        video_tensor = self.decode_latents(latents)\n",
						"\n",
						"        if output_type == \"pt\":\n",
						"            video = video_tensor\n",
						"        else:\n",
						"            video = tensor2vid(video_tensor)\n",
						"\n",
						"        if not return_dict:\n",
						"            return (video,)\n",
						"\n",
						"        return {\"frames\": video}\n",
						"\n",
						"    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt\n",
						"    def _encode_prompt(\n",
						"        self,\n",
						"        prompt,\n",
						"        num_images_per_prompt,\n",
						"        do_classifier_free_guidance,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Encodes the prompt into text encoder hidden states.\n",
						"\n",
						"        Args:\n",
						"             prompt (`str` or `List[str]`, *optional*):\n",
						"                prompt to be encoded\n",
						"            num_images_per_prompt (`int`):\n",
						"                number of images that should be generated per prompt\n",
						"            do_classifier_free_guidance (`bool`):\n",
						"                whether to use classifier free guidance or not\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"        \"\"\"\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        if prompt_embeds is None:\n",
						"            text_inputs = self.tokenizer(\n",
						"                prompt,\n",
						"                padding=\"max_length\",\n",
						"                max_length=self.tokenizer.model_max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"            text_input_ids = text_inputs.input_ids\n",
						"            untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n",
						"\n",
						"            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n",
						"                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])\n",
						"                print(\n",
						"                    \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n",
						"                    f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n",
						"                )\n",
						"\n",
						"            prompt_embeds = self.text_encoder(text_input_ids)\n",
						"            prompt_embeds = prompt_embeds[0]\n",
						"            prompt_embeds = torch.tensor(prompt_embeds)\n",
						"\n",
						"        bs_embed, seq_len, _ = prompt_embeds.shape\n",
						"        # duplicate text embeddings for each generation per prompt, using mps friendly method\n",
						"        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"        # get unconditional embeddings for classifier free guidance\n",
						"        if do_classifier_free_guidance and negative_prompt_embeds is None:\n",
						"            uncond_tokens: List[str]\n",
						"            if negative_prompt is None:\n",
						"                uncond_tokens = [\"\"] * batch_size\n",
						"            elif type(prompt) is not type(negative_prompt):\n",
						"                raise TypeError(f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\" f\" {type(prompt)}.\")\n",
						"            elif isinstance(negative_prompt, str):\n",
						"                uncond_tokens = [negative_prompt]\n",
						"            elif batch_size != len(negative_prompt):\n",
						"                raise ValueError(\n",
						"                    f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n",
						"                    f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n",
						"                    \" the batch size of `prompt`.\"\n",
						"                )\n",
						"            else:\n",
						"                uncond_tokens = negative_prompt\n",
						"\n",
						"            max_length = prompt_embeds.shape[1]\n",
						"            uncond_input = self.tokenizer(\n",
						"                uncond_tokens,\n",
						"                padding=\"max_length\",\n",
						"                max_length=max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"\n",
						"            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids)\n",
						"            negative_prompt_embeds = negative_prompt_embeds[0]\n",
						"            negative_prompt_embeds = torch.tensor(negative_prompt_embeds)\n",
						"\n",
						"        if do_classifier_free_guidance:\n",
						"            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n",
						"            seq_len = negative_prompt_embeds.shape[1]\n",
						"\n",
						"            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"            # For classifier free guidance, we need to do two forward passes.\n",
						"            # Here we concatenate the unconditional and text embeddings into a single batch\n",
						"            # to avoid doing two forward passes\n",
						"            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])\n",
						"\n",
						"        return prompt_embeds\n",
						"\n",
						"    def prepare_latents(\n",
						"        self,\n",
						"        batch_size,\n",
						"        num_channels_latents,\n",
						"        dtype,\n",
						"        generator,\n",
						"        latents=None,\n",
						"    ):\n",
						"        shape = (\n",
						"            batch_size,\n",
						"            num_channels_latents,\n",
						"            self.num_frames,\n",
						"            self.height // self.vae_scale_factor,\n",
						"            self.width // self.vae_scale_factor,\n",
						"        )\n",
						"        if isinstance(generator, list) and len(generator) != batch_size:\n",
						"            raise ValueError(\n",
						"                f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n",
						"                f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n",
						"            )\n",
						"\n",
						"        if latents is None:\n",
						"            latents = randn_tensor(shape, generator=generator, dtype=dtype)\n",
						"\n",
						"        # scale the initial noise by the standard deviation required by the scheduler\n",
						"        latents = latents * self.scheduler.init_noise_sigma\n",
						"        return latents\n",
						"\n",
						"    def check_inputs(\n",
						"        self,\n",
						"        prompt,\n",
						"        callback_steps,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds=None,\n",
						"        negative_prompt_embeds=None,\n",
						"    ):\n",
						"        if self.height % 8 != 0 or self.width % 8 != 0:\n",
						"            raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {self.height} and {self.width}.\")\n",
						"\n",
						"        if (callback_steps is None) or (callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)):\n",
						"            raise ValueError(f\"`callback_steps` has to be a positive integer but is {callback_steps} of type\" f\" {type(callback_steps)}.\")\n",
						"\n",
						"        if prompt is not None and prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to\" \" only forward one of the two.\"\n",
						"            )\n",
						"        elif prompt is None and prompt_embeds is None:\n",
						"            raise ValueError(\"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.\")\n",
						"        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):\n",
						"            raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n",
						"\n",
						"        if negative_prompt is not None and negative_prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:\"\n",
						"                f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n",
						"            )\n",
						"\n",
						"        if prompt_embeds is not None and negative_prompt_embeds is not None:\n",
						"            if prompt_embeds.shape != negative_prompt_embeds.shape:\n",
						"                raise ValueError(\n",
						"                    \"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but\"\n",
						"                    f\" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`\"\n",
						"                    f\" {negative_prompt_embeds.shape}.\"\n",
						"                )\n",
						"\n",
						"    def decode_latents(self, latents):\n",
						"        scale_factor = 0.18215\n",
						"        latents = 1 / scale_factor * latents\n",
						"\n",
						"        batch_size, channels, num_frames, height, width = latents.shape\n",
						"        latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)\n",
						"        image = self.vae_decoder(latents)[0]\n",
						"        image = torch.tensor(image)\n",
						"        video = (\n",
						"            image[None, :]\n",
						"            .reshape(\n",
						"                (\n",
						"                    batch_size,\n",
						"                    num_frames,\n",
						"                    -1,\n",
						"                )\n",
						"                + image.shape[2:]\n",
						"            )\n",
						"            .permute(0, 2, 1, 3, 4)\n",
						"        )\n",
						"        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16\n",
						"        video = video.float()\n",
						"        return video"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Inference with OpenVINO\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"from notebook_utils import device_widget\n",
						"\n",
						"device = device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_unet = core.compile_model(unet_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_vae_decoder = core.compile_model(vae_decoder_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_text_encoder = core.compile_model(text_encoder_xml, device_name=device.value)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Here we replace the pipeline parts with versions converted to OpenVINO IR and compiled to specific device. Note that we use original pipeline tokenizer and scheduler."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"ov_pipe = OVTextToVideoSDPipeline(ov_vae_decoder, ov_text_encoder, tokenizer, ov_unet, scheduler)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define a prompt\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"prompt = \"A panda eating bamboo on a rock.\""
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Let's generate a video for our prompt. For full list of arguments, see `__call__` function definition of `OVTextToVideoSDPipeline` class in [Build a pipeline](#Build-a-pipeline) section."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Video generation\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"frames = ov_pipe(prompt, num_inference_steps=25)[\"frames\"]"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"images[0].save(\"output.gif\", save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"with open(\"output.gif\", \"rb\") as gif_file:\n",
						"    b64 = f\"data:image/gif;base64,{base64.b64encode(gif_file.read()).decode()}\"\n",
						"IPython.display.HTML(f'<img src=\"{b64}\" />')"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Interactive demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def generate(prompt, seed, num_inference_steps, _=gr.Progress(track_tqdm=True)):\n",
						"    generator = torch.Generator().manual_seed(seed)\n",
						"    frames = ov_pipe(\n",
						"        prompt,\n",
						"        num_inference_steps=num_inference_steps,\n",
						"        generator=generator,\n",
						"    )[\"frames\"]\n",
						"    out_file = tempfile.NamedTemporaryFile(suffix=\".gif\", delete=False)\n",
						"    images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"    images[0].save(out_file, save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"    return out_file.name\n",
						"\n",
						"\n",
						"demo = gr.Interface(\n",
						"    generate,\n",
						"    [\n",
						"        gr.Textbox(label=\"Prompt\"),\n",
						"        gr.Slider(0, 1000000, value=42, label=\"Seed\", step=1),\n",
						"        gr.Slider(10, 50, value=25, label=\"Number of inference steps\", step=1),\n",
						"    ],\n",
						"    gr.Image(label=\"Result\"),\n",
						"    examples=[\n",
						"        [\"An astronaut riding a horse.\", 0, 25],\n",
						"        [\"A panda eating bamboo on a rock.\", 0, 25],\n",
						"        [\"Spiderman is surfing.\", 0, 25],\n",
						"    ],\n",
						"    allow_flagging=\"never\",\n",
						")\n",
						"\n",
						"try:\n",
						"    demo.queue().launch(debug=True)\n",
						"except Exception:\n",
						"    demo.queue().launch(share=True, debug=True)\n",
						"# if you are launching remotely, specify server_name and server_port\n",
						"# demo.launch(server_name='your server name', server_port='server port in int')\n",
						"# Read more in the docs: https://gradio.app/docs/"
					]
				}
			].map(fromJupyterCell)
			,
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Video generation with ZeroScope and OpenVINO\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Install and import required packages](#Install-and-import-required-packages)\n",
						"- [Load the model](#Load-the-model)\n",
						"- [Convert the model](#Convert-the-model)\n",
						"    - [Define the conversion function](#Define-the-conversion-function)\n",
						"    - [UNet](#UNet)\n",
						"    - [VAE](#VAE)\n",
						"    - [Text encoder](#Text-encoder)\n",
						"- [Build a pipeline](#Build-a-pipeline)\n",
						"- [Inference with OpenVINO](#Inference-with-OpenVINO)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Define a prompt](#Define-a-prompt)\n",
						"    - [Video generation](#Video-generation)\n",
						"- [Interactive demo](#Interactive-demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/zeroscope-text2video/zeroscope-text2video.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The ZeroScope model is a free and open-source text-to-video model that can generate realistic and engaging videos from text descriptions. It is based on the [Modelscope](https://modelscope.cn/models/damo/text-to-video-synthesis/summary) model, but it has been improved to produce higher-quality videos with a 16:9 aspect ratio and no Shutterstock watermark. The ZeroScope model is available in two versions: ZeroScope_v2 576w, which is optimized for rapid content creation at a resolution of 576x320 pixels, and ZeroScope_v2 XL, which upscales videos to a high-definition resolution of 1024x576.\n",
						"\n",
						"The ZeroScope model is trained on a dataset of over 9,000 videos and 29,000 tagged frames. It uses a diffusion model to generate videos, which means that it starts with a random noise image and gradually adds detail to it until it matches the text description. The ZeroScope model is still under development, but it has already been used to create some impressive videos. For example, it has been used to create videos of people dancing, playing sports, and even driving cars.\n",
						"\n",
						"The ZeroScope model is a powerful tool that can be used to create various videos, from simple animations to complex scenes. It is still under development, but it has the potential to revolutionize the way we create and consume video content.\n",
						"\n",
						"Both versions of the ZeroScope model are available on Hugging Face:\n",
						" - [ZeroScope_v2 576w](https://huggingface.co/cerspense/zeroscope_v2_576w)\n",
						" - [ZeroScope_v2 XL](https://huggingface.co/cerspense/zeroscope_v2_XL)\n",
						"\n",
						"We will use the first one."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"<div class=\"alert alert-block alert-warning\">\n",
						"    This tutorial requires at least 24GB of free memory to generate a video with a frame size of 432x240 and 16 frames. Increasing either of these values will require more memory and take more time.\n",
						"</div>"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Install and import required packages\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"To work with text-to-video synthesis model, we will use Hugging Face's [Diffusers](https://github.com/huggingface/diffusers) library. It provides already pretrained model from `cerspense`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu  \"diffusers>=0.18.0\" \"torch>=2.1\" transformers \"openvino>=2023.1.0\" numpy \"gradio>=4.19\""
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import gc\n",
						"from typing import Optional, Union, List, Callable\n",
						"import base64\n",
						"import tempfile\n",
						"import warnings\n",
						"\n",
						"import diffusers\n",
						"import transformers\n",
						"import numpy as np\n",
						"import IPython\n",
						"import torch\n",
						"import PIL\n",
						"import gradio as gr\n",
						"\n",
						"import openvino as ov"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Original 576x320 inference requires a lot of RAM (>100GB), so let's run our example on a smaller frame size, keeping the same aspect ratio. Try reducing values below to reduce the memory consumption."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"WIDTH = 432  # must be divisible by 8\n",
						"HEIGHT = 240  # must be divisible by 8\n",
						"NUM_FRAMES = 16"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The model is loaded from HuggingFace using `.from_pretrained` method of `diffusers.DiffusionPipeline`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"pipe = diffusers.DiffusionPipeline.from_pretrained(\"cerspense/zeroscope_v2_576w\")"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet = pipe.unet\n",
						"unet.eval()\n",
						"vae = pipe.vae\n",
						"vae.eval()\n",
						"text_encoder = pipe.text_encoder\n",
						"text_encoder.eval()\n",
						"tokenizer = pipe.tokenizer\n",
						"scheduler = pipe.scheduler\n",
						"vae_scale_factor = pipe.vae_scale_factor\n",
						"unet_in_channels = pipe.unet.config.in_channels\n",
						"sample_width = WIDTH // vae_scale_factor\n",
						"sample_height = HEIGHT // vae_scale_factor\n",
						"del pipe\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Convert the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The architecture for generating videos from text comprises three distinct sub-networks: one for extracting text features, another for translating text features into the video latent space using a diffusion model, and a final one for mapping the video latent space to the visual space. The collective parameters of the entire model amount to approximately 1.7 billion. It's capable of processing English input. The diffusion model is built upon the Unet3D model and achieves video generation by iteratively denoising a starting point of pure Gaussian noise video."
					]
				},
				{
					"cell_type": "markdown",
					"source": []
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define the conversion function\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Model components are PyTorch modules, that can be converted with `ov.convert_model` function directly. We also use `ov.save_model` function to serialize the result of conversion."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"\n",
						"def convert(model: torch.nn.Module, xml_path: str, **convert_kwargs) -> Path:\n",
						"    xml_path = Path(xml_path)\n",
						"    if not xml_path.exists():\n",
						"        xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
						"        with torch.no_grad():\n",
						"            converted_model = ov.convert_model(model, **convert_kwargs)\n",
						"        ov.save_model(converted_model, xml_path)\n",
						"        del converted_model\n",
						"        gc.collect()\n",
						"        torch._C._jit_clear_class_registry()\n",
						"        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
						"        torch.jit._state._clear_class_state()\n",
						"    return xml_path"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### UNet\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text-to-video generation pipeline main component is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet_xml_path = convert(\n",
						"    unet,\n",
						"    \"models/unet.xml\",\n",
						"    example_input={\n",
						"        \"sample\": torch.randn(2, 4, 2, int(sample_height // 2), int(sample_width // 2)),\n",
						"        \"timestep\": torch.tensor(1),\n",
						"        \"encoder_hidden_states\": torch.randn(2, 77, 1024),\n",
						"    },\n",
						"    input=[\n",
						"        (\"sample\", (2, 4, NUM_FRAMES, sample_height, sample_width)),\n",
						"        (\"timestep\", ()),\n",
						"        (\"encoder_hidden_states\", (2, 77, 1024)),\n",
						"    ],\n",
						")\n",
						"del unet\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### VAE\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Variational autoencoder (VAE) uses UNet output to decode latents to visual representations. Our VAE model has KL loss for encoding images into latents and decoding latent representations into images. For inference, we need only decoder part."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"class VaeDecoderWrapper(torch.nn.Module):\n",
						"    def __init__(self, vae):\n",
						"        super().__init__()\n",
						"        self.vae = vae\n",
						"\n",
						"    def forward(self, z: torch.FloatTensor):\n",
						"        return self.vae.decode(z)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"vae_decoder_xml_path = convert(\n",
						"    VaeDecoderWrapper(vae),\n",
						"    \"models/vae.xml\",\n",
						"    example_input=torch.randn(2, 4, 32, 32),\n",
						"    input=((NUM_FRAMES, 4, sample_height, sample_width)),\n",
						")\n",
						"del vae\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Text encoder\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text encoder is used to encode the input prompt to tensor. Default tensor length is 77."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"text_encoder_xml = convert(\n",
						"    text_encoder,\n",
						"    \"models/text_encoder.xml\",\n",
						"    example_input=torch.ones(1, 77, dtype=torch.int64),\n",
						"    input=((1, 77), ov.Type.i64),\n",
						")\n",
						"del text_encoder\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Build a pipeline\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> List[np.ndarray]:\n",
						"    # This code is copied from https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78\n",
						"    # reshape to ncfhw\n",
						"    mean = torch.tensor(mean, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    std = torch.tensor(std, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    # unnormalize back to [0,1]\n",
						"    video = video.mul_(std).add_(mean)\n",
						"    video.clamp_(0, 1)\n",
						"    # prepare the final outputs\n",
						"    i, c, f, h, w = video.shape\n",
						"    images = video.permute(2, 3, 0, 4, 1).reshape(f, h, i * w, c)  # 1st (frames, h, batch_size, w, c) 2nd (frames, h, batch_size * w, c)\n",
						"    images = images.unbind(dim=0)  # prepare a list of indvidual (consecutive frames)\n",
						"    images = [(image.cpu().numpy() * 255).astype(\"uint8\") for image in images]  # f h w c\n",
						"    return images"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"try:\n",
						"    from diffusers.utils import randn_tensor\n",
						"except ImportError:\n",
						"    from diffusers.utils.torch_utils import randn_tensor\n",
						"\n",
						"\n",
						"class OVTextToVideoSDPipeline(diffusers.DiffusionPipeline):\n",
						"    def __init__(\n",
						"        self,\n",
						"        vae_decoder: ov.CompiledModel,\n",
						"        text_encoder: ov.CompiledModel,\n",
						"        tokenizer: transformers.CLIPTokenizer,\n",
						"        unet: ov.CompiledModel,\n",
						"        scheduler: diffusers.schedulers.DDIMScheduler,\n",
						"    ):\n",
						"        super().__init__()\n",
						"\n",
						"        self.vae_decoder = vae_decoder\n",
						"        self.text_encoder = text_encoder\n",
						"        self.tokenizer = tokenizer\n",
						"        self.unet = unet\n",
						"        self.scheduler = scheduler\n",
						"        self.vae_scale_factor = vae_scale_factor\n",
						"        self.unet_in_channels = unet_in_channels\n",
						"        self.width = WIDTH\n",
						"        self.height = HEIGHT\n",
						"        self.num_frames = NUM_FRAMES\n",
						"\n",
						"    def __call__(\n",
						"        self,\n",
						"        prompt: Union[str, List[str]] = None,\n",
						"        num_inference_steps: int = 50,\n",
						"        guidance_scale: float = 9.0,\n",
						"        negative_prompt: Optional[Union[str, List[str]]] = None,\n",
						"        eta: float = 0.0,\n",
						"        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n",
						"        latents: Optional[torch.FloatTensor] = None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        output_type: Optional[str] = \"np\",\n",
						"        return_dict: bool = True,\n",
						"        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n",
						"        callback_steps: int = 1,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Function invoked when calling the pipeline for generation.\n",
						"\n",
						"        Args:\n",
						"            prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.\n",
						"                instead.\n",
						"            num_inference_steps (`int`, *optional*, defaults to 50):\n",
						"                The number of denoising steps. More denoising steps usually lead to a higher quality videos at the\n",
						"                expense of slower inference.\n",
						"            guidance_scale (`float`, *optional*, defaults to 7.5):\n",
						"                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).\n",
						"                `guidance_scale` is defined as `w` of equation 2. of [Imagen\n",
						"                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n",
						"                1`. Higher guidance scale encourages to generate videos that are closely linked to the text `prompt`,\n",
						"                usually at the expense of lower video quality.\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the video generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            eta (`float`, *optional*, defaults to 0.0):\n",
						"                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n",
						"                [`schedulers.DDIMScheduler`], will be ignored for others.\n",
						"            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n",
						"                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n",
						"                to make generation deterministic.\n",
						"            latents (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video\n",
						"                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n",
						"                tensor will ge generated by sampling using the supplied random `generator`. Latents should be of shape\n",
						"                `(batch_size, num_channel, num_frames, height, width)`.\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"            output_type (`str`, *optional*, defaults to `\"np\"`):\n",
						"                The output format of the generate video. Choose between `torch.FloatTensor` or `np.array`.\n",
						"            return_dict (`bool`, *optional*, defaults to `True`):\n",
						"                Whether or not to return a [`~pipelines.stable_diffusion.TextToVideoSDPipelineOutput`] instead of a\n",
						"                plain tuple.\n",
						"            callback (`Callable`, *optional*):\n",
						"                A function that will be called every `callback_steps` steps during inference. The function will be\n",
						"                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n",
						"            callback_steps (`int`, *optional*, defaults to 1):\n",
						"                The frequency at which the `callback` function will be called. If not specified, the callback will be\n",
						"                called at every step.\n",
						"\n",
						"        Returns:\n",
						"            `List[np.ndarray]`: generated video frames\n",
						"        \"\"\"\n",
						"\n",
						"        num_images_per_prompt = 1\n",
						"\n",
						"        # 1. Check inputs. Raise error if not correct\n",
						"        self.check_inputs(\n",
						"            prompt,\n",
						"            callback_steps,\n",
						"            negative_prompt,\n",
						"            prompt_embeds,\n",
						"            negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 2. Define call parameters\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n",
						"        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n",
						"        # corresponds to doing no classifier free guidance.\n",
						"        do_classifier_free_guidance = guidance_scale > 1.0\n",
						"\n",
						"        # 3. Encode input prompt\n",
						"        prompt_embeds = self._encode_prompt(\n",
						"            prompt,\n",
						"            num_images_per_prompt,\n",
						"            do_classifier_free_guidance,\n",
						"            negative_prompt,\n",
						"            prompt_embeds=prompt_embeds,\n",
						"            negative_prompt_embeds=negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 4. Prepare timesteps\n",
						"        self.scheduler.set_timesteps(num_inference_steps)\n",
						"        timesteps = self.scheduler.timesteps\n",
						"\n",
						"        # 5. Prepare latent variables\n",
						"        num_channels_latents = self.unet_in_channels\n",
						"        latents = self.prepare_latents(\n",
						"            batch_size * num_images_per_prompt,\n",
						"            num_channels_latents,\n",
						"            prompt_embeds.dtype,\n",
						"            generator,\n",
						"            latents,\n",
						"        )\n",
						"\n",
						"        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline\n",
						"        extra_step_kwargs = {\"generator\": generator, \"eta\": eta}\n",
						"\n",
						"        # 7. Denoising loop\n",
						"        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n",
						"        with self.progress_bar(total=num_inference_steps) as progress_bar:\n",
						"            for i, t in enumerate(timesteps):\n",
						"                # expand the latents if we are doing classifier free guidance\n",
						"                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n",
						"                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
						"\n",
						"                # predict the noise residual\n",
						"                noise_pred = self.unet(\n",
						"                    {\n",
						"                        \"sample\": latent_model_input,\n",
						"                        \"timestep\": t,\n",
						"                        \"encoder_hidden_states\": prompt_embeds,\n",
						"                    }\n",
						"                )[0]\n",
						"                noise_pred = torch.tensor(noise_pred)\n",
						"\n",
						"                # perform guidance\n",
						"                if do_classifier_free_guidance:\n",
						"                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
						"                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
						"\n",
						"                # reshape latents\n",
						"                bsz, channel, frames, width, height = latents.shape\n",
						"                latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"                noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"\n",
						"                # compute the previous noisy sample x_t -> x_t-1\n",
						"                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample\n",
						"\n",
						"                # reshape latents back\n",
						"                latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4)\n",
						"\n",
						"                # call the callback, if provided\n",
						"                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n",
						"                    progress_bar.update()\n",
						"                    if callback is not None and i % callback_steps == 0:\n",
						"                        callback(i, t, latents)\n",
						"\n",
						"        video_tensor = self.decode_latents(latents)\n",
						"\n",
						"        if output_type == \"pt\":\n",
						"            video = video_tensor\n",
						"        else:\n",
						"            video = tensor2vid(video_tensor)\n",
						"\n",
						"        if not return_dict:\n",
						"            return (video,)\n",
						"\n",
						"        return {\"frames\": video}\n",
						"\n",
						"    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt\n",
						"    def _encode_prompt(\n",
						"        self,\n",
						"        prompt,\n",
						"        num_images_per_prompt,\n",
						"        do_classifier_free_guidance,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Encodes the prompt into text encoder hidden states.\n",
						"\n",
						"        Args:\n",
						"             prompt (`str` or `List[str]`, *optional*):\n",
						"                prompt to be encoded\n",
						"            num_images_per_prompt (`int`):\n",
						"                number of images that should be generated per prompt\n",
						"            do_classifier_free_guidance (`bool`):\n",
						"                whether to use classifier free guidance or not\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"        \"\"\"\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        if prompt_embeds is None:\n",
						"            text_inputs = self.tokenizer(\n",
						"                prompt,\n",
						"                padding=\"max_length\",\n",
						"                max_length=self.tokenizer.model_max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"            text_input_ids = text_inputs.input_ids\n",
						"            untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n",
						"\n",
						"            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n",
						"                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])\n",
						"                print(\n",
						"                    \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n",
						"                    f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n",
						"                )\n",
						"\n",
						"            prompt_embeds = self.text_encoder(text_input_ids)\n",
						"            prompt_embeds = prompt_embeds[0]\n",
						"            prompt_embeds = torch.tensor(prompt_embeds)\n",
						"\n",
						"        bs_embed, seq_len, _ = prompt_embeds.shape\n",
						"        # duplicate text embeddings for each generation per prompt, using mps friendly method\n",
						"        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"        # get unconditional embeddings for classifier free guidance\n",
						"        if do_classifier_free_guidance and negative_prompt_embeds is None:\n",
						"            uncond_tokens: List[str]\n",
						"            if negative_prompt is None:\n",
						"                uncond_tokens = [\"\"] * batch_size\n",
						"            elif type(prompt) is not type(negative_prompt):\n",
						"                raise TypeError(f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\" f\" {type(prompt)}.\")\n",
						"            elif isinstance(negative_prompt, str):\n",
						"                uncond_tokens = [negative_prompt]\n",
						"            elif batch_size != len(negative_prompt):\n",
						"                raise ValueError(\n",
						"                    f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n",
						"                    f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n",
						"                    \" the batch size of `prompt`.\"\n",
						"                )\n",
						"            else:\n",
						"                uncond_tokens = negative_prompt\n",
						"\n",
						"            max_length = prompt_embeds.shape[1]\n",
						"            uncond_input = self.tokenizer(\n",
						"                uncond_tokens,\n",
						"                padding=\"max_length\",\n",
						"                max_length=max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"\n",
						"            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids)\n",
						"            negative_prompt_embeds = negative_prompt_embeds[0]\n",
						"            negative_prompt_embeds = torch.tensor(negative_prompt_embeds)\n",
						"\n",
						"        if do_classifier_free_guidance:\n",
						"            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n",
						"            seq_len = negative_prompt_embeds.shape[1]\n",
						"\n",
						"            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"            # For classifier free guidance, we need to do two forward passes.\n",
						"            # Here we concatenate the unconditional and text embeddings into a single batch\n",
						"            # to avoid doing two forward passes\n",
						"            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])\n",
						"\n",
						"        return prompt_embeds\n",
						"\n",
						"    def prepare_latents(\n",
						"        self,\n",
						"        batch_size,\n",
						"        num_channels_latents,\n",
						"        dtype,\n",
						"        generator,\n",
						"        latents=None,\n",
						"    ):\n",
						"        shape = (\n",
						"            batch_size,\n",
						"            num_channels_latents,\n",
						"            self.num_frames,\n",
						"            self.height // self.vae_scale_factor,\n",
						"            self.width // self.vae_scale_factor,\n",
						"        )\n",
						"        if isinstance(generator, list) and len(generator) != batch_size:\n",
						"            raise ValueError(\n",
						"                f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n",
						"                f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n",
						"            )\n",
						"\n",
						"        if latents is None:\n",
						"            latents = randn_tensor(shape, generator=generator, dtype=dtype)\n",
						"\n",
						"        # scale the initial noise by the standard deviation required by the scheduler\n",
						"        latents = latents * self.scheduler.init_noise_sigma\n",
						"        return latents\n",
						"\n",
						"    def check_inputs(\n",
						"        self,\n",
						"        prompt,\n",
						"        callback_steps,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds=None,\n",
						"        negative_prompt_embeds=None,\n",
						"    ):\n",
						"        if self.height % 8 != 0 or self.width % 8 != 0:\n",
						"            raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {self.height} and {self.width}.\")\n",
						"\n",
						"        if (callback_steps is None) or (callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)):\n",
						"            raise ValueError(f\"`callback_steps` has to be a positive integer but is {callback_steps} of type\" f\" {type(callback_steps)}.\")\n",
						"\n",
						"        if prompt is not None and prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to\" \" only forward one of the two.\"\n",
						"            )\n",
						"        elif prompt is None and prompt_embeds is None:\n",
						"            raise ValueError(\"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.\")\n",
						"        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):\n",
						"            raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n",
						"\n",
						"        if negative_prompt is not None and negative_prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:\"\n",
						"                f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n",
						"            )\n",
						"\n",
						"        if prompt_embeds is not None and negative_prompt_embeds is not None:\n",
						"            if prompt_embeds.shape != negative_prompt_embeds.shape:\n",
						"                raise ValueError(\n",
						"                    \"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but\"\n",
						"                    f\" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`\"\n",
						"                    f\" {negative_prompt_embeds.shape}.\"\n",
						"                )\n",
						"\n",
						"    def decode_latents(self, latents):\n",
						"        scale_factor = 0.18215\n",
						"        latents = 1 / scale_factor * latents\n",
						"\n",
						"        batch_size, channels, num_frames, height, width = latents.shape\n",
						"        latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)\n",
						"        image = self.vae_decoder(latents)[0]\n",
						"        image = torch.tensor(image)\n",
						"        video = (\n",
						"            image[None, :]\n",
						"            .reshape(\n",
						"                (\n",
						"                    batch_size,\n",
						"                    num_frames,\n",
						"                    -1,\n",
						"                )\n",
						"                + image.shape[2:]\n",
						"            )\n",
						"            .permute(0, 2, 1, 3, 4)\n",
						"        )\n",
						"        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16\n",
						"        video = video.float()\n",
						"        return video"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Inference with OpenVINO\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"from notebook_utils import device_widget\n",
						"\n",
						"device = device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_unet = core.compile_model(unet_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_vae_decoder = core.compile_model(vae_decoder_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_text_encoder = core.compile_model(text_encoder_xml, device_name=device.value)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Here we replace the pipeline parts with versions converted to OpenVINO IR and compiled to specific device. Note that we use original pipeline tokenizer and scheduler."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"ov_pipe = OVTextToVideoSDPipeline(ov_vae_decoder, ov_text_encoder, tokenizer, ov_unet, scheduler)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define a prompt\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"prompt = \"A panda eating bamboo on a rock.\""
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Let's generate a video for our prompt. For full list of arguments, see `__call__` function definition of `OVTextToVideoSDPipeline` class in [Build a pipeline](#Build-a-pipeline) section."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Video generation\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"frames = ov_pipe(prompt, num_inference_steps=25)[\"frames\"]"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"images[0].save(\"output.gif\", save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"with open(\"output.gif\", \"rb\") as gif_file:\n",
						"    b64 = f\"data:image/gif;base64,{base64.b64encode(gif_file.read()).decode()}\"\n",
						"IPython.display.HTML(f'<img src=\"{b64}\" />')"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Interactive demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"test_replace": {
							"    demo.queue().launch(debug=True)": "    demo.queue.launch()",
							"    demo.queue().launch(share=True, debug=True)": "    demo.queue().launch(share=True)"
						}
					},
					"source": [
						"def generate(prompt, seed, num_inference_steps, _=gr.Progress(track_tqdm=True)):\n",
						"    generator = torch.Generator().manual_seed(seed)\n",
						"    frames = ov_pipe(\n",
						"        prompt,\n",
						"        num_inference_steps=num_inference_steps,\n",
						"        generator=generator,\n",
						"    )[\"frames\"]\n",
						"    out_file = tempfile.NamedTemporaryFile(suffix=\".gif\", delete=False)\n",
						"    images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"    images[0].save(out_file, save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"    return out_file.name"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"outputs": [],
					"source": [
						"if not Path(\"gradio_helper.py\").exists():\n",
						"    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/zeroscope-text2video/gradio_helper.py\")\n",
						"    open(\"gradio_helper.py\", \"w\").write(r.text)\n",
						"\n",
						"from gradio_helper import make_demo\n",
						"\n",
						"demo = make_demo(fn=generate)\n",
						"\n",
						"try:\n",
						"    demo.queue().launch(debug=True)\n",
						"except Exception:\n",
						"    demo.queue().launch(share=True, debug=True)\n",
						"# If you are launching remotely, specify server_name and server_port\n",
						"# EXAMPLE: `demo.launch(server_name='your server name', server_port='server port in int')`\n",
						"# To learn more please refer to the Gradio docs: https://gradio.app/docs/"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"# please uncomment and run this cell for stopping gradio interface\n",
						"# demo.close()"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
			{ modified: 24, original: 24 },
			{ modified: 25, original: 25 },
			{ modified: 26, original: 26 },
			{ modified: 27, original: 27 },
			{ modified: 28, original: 28 },
			{ modified: 29, original: 29 },
			{ modified: 30, original: 30 },
			{ modified: 31, original: 31 },
			{ modified: 32, original: 32 },
			{ modified: 33, original: 33 },
			{ modified: 34, original: 34 },
			{ modified: 35, original: 35 },
			{ modified: 36, original: 36 },
			{ modified: 37, original: 37 },
			{ modified: 38, original: 38 },
			{ modified: 39, original: 39 },
			{ modified: 40, original: 40 },
			{ modified: 41, original: 41 },
			{ modified: 42, original: 42 },
			{ modified: 43, original: 43 },
			{ modified: 44, original: 44 },
			{ modified: 45, original: 45 },
			{ modified: 46, original: 46 },
			{ modified: 47, original: 47 },
			{ modified: 48, original: 48 },
			{ modified: 49, original: 49 },
			{ modified: 50, original: -1 },
			{ modified: 51, original: -1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 49, originalLength: 1, modifiedStart: 49, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 50, originalLength: 0, modifiedStart: 50, modifiedLength: 2 } satisfies IDiffChange,
		]);
	});
	test('Modification of three cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Video generation with ZeroScope and OpenVINO\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Install and import required packages](#Install-and-import-required-packages)\n",
						"- [Load the model](#Load-the-model)\n",
						"- [Convert the model](#Convert-the-model)\n",
						"    - [Define the conversion function](#Define-the-conversion-function)\n",
						"    - [UNet](#UNet)\n",
						"    - [VAE](#VAE)\n",
						"    - [Text encoder](#Text-encoder)\n",
						"- [Build a pipeline](#Build-a-pipeline)\n",
						"- [Inference with OpenVINO](#Inference-with-OpenVINO)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Define a prompt](#Define-a-prompt)\n",
						"    - [Video generation](#Video-generation)\n",
						"- [Interactive demo](#Interactive-demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/zeroscope-text2video/zeroscope-text2video.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The ZeroScope model is a free and open-source text-to-video model that can generate realistic and engaging videos from text descriptions. It is based on the [Modelscope](https://modelscope.cn/models/damo/text-to-video-synthesis/summary) model, but it has been improved to produce higher-quality videos with a 16:9 aspect ratio and no Shutterstock watermark. The ZeroScope model is available in two versions: ZeroScope_v2 576w, which is optimized for rapid content creation at a resolution of 576x320 pixels, and ZeroScope_v2 XL, which upscales videos to a high-definition resolution of 1024x576.\n",
						"\n",
						"The ZeroScope model is trained on a dataset of over 9,000 videos and 29,000 tagged frames. It uses a diffusion model to generate videos, which means that it starts with a random noise image and gradually adds detail to it until it matches the text description. The ZeroScope model is still under development, but it has already been used to create some impressive videos. For example, it has been used to create videos of people dancing, playing sports, and even driving cars.\n",
						"\n",
						"The ZeroScope model is a powerful tool that can be used to create various videos, from simple animations to complex scenes. It is still under development, but it has the potential to revolutionize the way we create and consume video content.\n",
						"\n",
						"Both versions of the ZeroScope model are available on Hugging Face:\n",
						" - [ZeroScope_v2 576w](https://huggingface.co/cerspense/zeroscope_v2_576w)\n",
						" - [ZeroScope_v2 XL](https://huggingface.co/cerspense/zeroscope_v2_XL)\n",
						"\n",
						"We will use the first one."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"<div class=\"alert alert-block alert-warning\">\n",
						"    This tutorial requires at least 24GB of free memory to generate a video with a frame size of 432x240 and 16 frames. Increasing either of these values will require more memory and take more time.\n",
						"</div>"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Install and import required packages\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"To work with text-to-video synthesis model, we will use Hugging Face's [Diffusers](https://github.com/huggingface/diffusers) library. It provides already pretrained model from `cerspense`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu  \"diffusers>=0.18.0\" \"torch>=2.1\" transformers \"openvino>=2023.1.0\" numpy \"gradio>=4.19\""
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import gc\n",
						"from typing import Optional, Union, List, Callable\n",
						"import base64\n",
						"import tempfile\n",
						"import warnings\n",
						"\n",
						"import diffusers\n",
						"import transformers\n",
						"import numpy as np\n",
						"import IPython\n",
						"import ipywidgets as widgets\n",
						"import torch\n",
						"import PIL\n",
						"import gradio as gr\n",
						"\n",
						"import openvino as ov"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Original 576x320 inference requires a lot of RAM (>100GB), so let's run our example on a smaller frame size, keeping the same aspect ratio. Try reducing values below to reduce the memory consumption."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"WIDTH = 432  # must be divisible by 8\n",
						"HEIGHT = 240  # must be divisible by 8\n",
						"NUM_FRAMES = 16"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The model is loaded from HuggingFace using `.from_pretrained` method of `diffusers.DiffusionPipeline`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"pipe = diffusers.DiffusionPipeline.from_pretrained(\"cerspense/zeroscope_v2_576w\")"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet = pipe.unet\n",
						"unet.eval()\n",
						"vae = pipe.vae\n",
						"vae.eval()\n",
						"text_encoder = pipe.text_encoder\n",
						"text_encoder.eval()\n",
						"tokenizer = pipe.tokenizer\n",
						"scheduler = pipe.scheduler\n",
						"vae_scale_factor = pipe.vae_scale_factor\n",
						"unet_in_channels = pipe.unet.config.in_channels\n",
						"sample_width = WIDTH // vae_scale_factor\n",
						"sample_height = HEIGHT // vae_scale_factor\n",
						"del pipe\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Convert the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The architecture for generating videos from text comprises three distinct sub-networks: one for extracting text features, another for translating text features into the video latent space using a diffusion model, and a final one for mapping the video latent space to the visual space. The collective parameters of the entire model amount to approximately 1.7 billion. It's capable of processing English input. The diffusion model is built upon the Unet3D model and achieves video generation by iteratively denoising a starting point of pure Gaussian noise video."
					]
				},
				{
					"cell_type": "markdown",
					"source": []
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define the conversion function\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Model components are PyTorch modules, that can be converted with `ov.convert_model` function directly. We also use `ov.save_model` function to serialize the result of conversion."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"\n",
						"def convert(model: torch.nn.Module, xml_path: str, **convert_kwargs) -> Path:\n",
						"    xml_path = Path(xml_path)\n",
						"    if not xml_path.exists():\n",
						"        xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
						"        with torch.no_grad():\n",
						"            converted_model = ov.convert_model(model, **convert_kwargs)\n",
						"        ov.save_model(converted_model, xml_path)\n",
						"        del converted_model\n",
						"        gc.collect()\n",
						"        torch._C._jit_clear_class_registry()\n",
						"        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
						"        torch.jit._state._clear_class_state()\n",
						"    return xml_path"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### UNet\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text-to-video generation pipeline main component is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet_xml_path = convert(\n",
						"    unet,\n",
						"    \"models/unet.xml\",\n",
						"    example_input={\n",
						"        \"sample\": torch.randn(2, 4, 2, int(sample_height // 2), int(sample_width // 2)),\n",
						"        \"timestep\": torch.tensor(1),\n",
						"        \"encoder_hidden_states\": torch.randn(2, 77, 1024),\n",
						"    },\n",
						"    input=[\n",
						"        (\"sample\", (2, 4, NUM_FRAMES, sample_height, sample_width)),\n",
						"        (\"timestep\", ()),\n",
						"        (\"encoder_hidden_states\", (2, 77, 1024)),\n",
						"    ],\n",
						")\n",
						"del unet\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### VAE\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Variational autoencoder (VAE) uses UNet output to decode latents to visual representations. Our VAE model has KL loss for encoding images into latents and decoding latent representations into images. For inference, we need only decoder part."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"class VaeDecoderWrapper(torch.nn.Module):\n",
						"    def __init__(self, vae):\n",
						"        super().__init__()\n",
						"        self.vae = vae\n",
						"\n",
						"    def forward(self, z: torch.FloatTensor):\n",
						"        return self.vae.decode(z)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"vae_decoder_xml_path = convert(\n",
						"    VaeDecoderWrapper(vae),\n",
						"    \"models/vae.xml\",\n",
						"    example_input=torch.randn(2, 4, 32, 32),\n",
						"    input=((NUM_FRAMES, 4, sample_height, sample_width)),\n",
						")\n",
						"del vae\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Text encoder\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text encoder is used to encode the input prompt to tensor. Default tensor length is 77."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"text_encoder_xml = convert(\n",
						"    text_encoder,\n",
						"    \"models/text_encoder.xml\",\n",
						"    example_input=torch.ones(1, 77, dtype=torch.int64),\n",
						"    input=((1, 77), ov.Type.i64),\n",
						")\n",
						"del text_encoder\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Build a pipeline\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> List[np.ndarray]:\n",
						"    # This code is copied from https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78\n",
						"    # reshape to ncfhw\n",
						"    mean = torch.tensor(mean, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    std = torch.tensor(std, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    # unnormalize back to [0,1]\n",
						"    video = video.mul_(std).add_(mean)\n",
						"    video.clamp_(0, 1)\n",
						"    # prepare the final outputs\n",
						"    i, c, f, h, w = video.shape\n",
						"    images = video.permute(2, 3, 0, 4, 1).reshape(f, h, i * w, c)  # 1st (frames, h, batch_size, w, c) 2nd (frames, h, batch_size * w, c)\n",
						"    images = images.unbind(dim=0)  # prepare a list of indvidual (consecutive frames)\n",
						"    images = [(image.cpu().numpy() * 255).astype(\"uint8\") for image in images]  # f h w c\n",
						"    return images"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"try:\n",
						"    from diffusers.utils import randn_tensor\n",
						"except ImportError:\n",
						"    from diffusers.utils.torch_utils import randn_tensor\n",
						"\n",
						"\n",
						"class OVTextToVideoSDPipeline(diffusers.DiffusionPipeline):\n",
						"    def __init__(\n",
						"        self,\n",
						"        vae_decoder: ov.CompiledModel,\n",
						"        text_encoder: ov.CompiledModel,\n",
						"        tokenizer: transformers.CLIPTokenizer,\n",
						"        unet: ov.CompiledModel,\n",
						"        scheduler: diffusers.schedulers.DDIMScheduler,\n",
						"    ):\n",
						"        super().__init__()\n",
						"\n",
						"        self.vae_decoder = vae_decoder\n",
						"        self.text_encoder = text_encoder\n",
						"        self.tokenizer = tokenizer\n",
						"        self.unet = unet\n",
						"        self.scheduler = scheduler\n",
						"        self.vae_scale_factor = vae_scale_factor\n",
						"        self.unet_in_channels = unet_in_channels\n",
						"        self.width = WIDTH\n",
						"        self.height = HEIGHT\n",
						"        self.num_frames = NUM_FRAMES\n",
						"\n",
						"    def __call__(\n",
						"        self,\n",
						"        prompt: Union[str, List[str]] = None,\n",
						"        num_inference_steps: int = 50,\n",
						"        guidance_scale: float = 9.0,\n",
						"        negative_prompt: Optional[Union[str, List[str]]] = None,\n",
						"        eta: float = 0.0,\n",
						"        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n",
						"        latents: Optional[torch.FloatTensor] = None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        output_type: Optional[str] = \"np\",\n",
						"        return_dict: bool = True,\n",
						"        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n",
						"        callback_steps: int = 1,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Function invoked when calling the pipeline for generation.\n",
						"\n",
						"        Args:\n",
						"            prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.\n",
						"                instead.\n",
						"            num_inference_steps (`int`, *optional*, defaults to 50):\n",
						"                The number of denoising steps. More denoising steps usually lead to a higher quality videos at the\n",
						"                expense of slower inference.\n",
						"            guidance_scale (`float`, *optional*, defaults to 7.5):\n",
						"                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).\n",
						"                `guidance_scale` is defined as `w` of equation 2. of [Imagen\n",
						"                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n",
						"                1`. Higher guidance scale encourages to generate videos that are closely linked to the text `prompt`,\n",
						"                usually at the expense of lower video quality.\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the video generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            eta (`float`, *optional*, defaults to 0.0):\n",
						"                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n",
						"                [`schedulers.DDIMScheduler`], will be ignored for others.\n",
						"            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n",
						"                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n",
						"                to make generation deterministic.\n",
						"            latents (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video\n",
						"                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n",
						"                tensor will ge generated by sampling using the supplied random `generator`. Latents should be of shape\n",
						"                `(batch_size, num_channel, num_frames, height, width)`.\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"            output_type (`str`, *optional*, defaults to `\"np\"`):\n",
						"                The output format of the generate video. Choose between `torch.FloatTensor` or `np.array`.\n",
						"            return_dict (`bool`, *optional*, defaults to `True`):\n",
						"                Whether or not to return a [`~pipelines.stable_diffusion.TextToVideoSDPipelineOutput`] instead of a\n",
						"                plain tuple.\n",
						"            callback (`Callable`, *optional*):\n",
						"                A function that will be called every `callback_steps` steps during inference. The function will be\n",
						"                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n",
						"            callback_steps (`int`, *optional*, defaults to 1):\n",
						"                The frequency at which the `callback` function will be called. If not specified, the callback will be\n",
						"                called at every step.\n",
						"\n",
						"        Returns:\n",
						"            `List[np.ndarray]`: generated video frames\n",
						"        \"\"\"\n",
						"\n",
						"        num_images_per_prompt = 1\n",
						"\n",
						"        # 1. Check inputs. Raise error if not correct\n",
						"        self.check_inputs(\n",
						"            prompt,\n",
						"            callback_steps,\n",
						"            negative_prompt,\n",
						"            prompt_embeds,\n",
						"            negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 2. Define call parameters\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n",
						"        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n",
						"        # corresponds to doing no classifier free guidance.\n",
						"        do_classifier_free_guidance = guidance_scale > 1.0\n",
						"\n",
						"        # 3. Encode input prompt\n",
						"        prompt_embeds = self._encode_prompt(\n",
						"            prompt,\n",
						"            num_images_per_prompt,\n",
						"            do_classifier_free_guidance,\n",
						"            negative_prompt,\n",
						"            prompt_embeds=prompt_embeds,\n",
						"            negative_prompt_embeds=negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 4. Prepare timesteps\n",
						"        self.scheduler.set_timesteps(num_inference_steps)\n",
						"        timesteps = self.scheduler.timesteps\n",
						"\n",
						"        # 5. Prepare latent variables\n",
						"        num_channels_latents = self.unet_in_channels\n",
						"        latents = self.prepare_latents(\n",
						"            batch_size * num_images_per_prompt,\n",
						"            num_channels_latents,\n",
						"            prompt_embeds.dtype,\n",
						"            generator,\n",
						"            latents,\n",
						"        )\n",
						"\n",
						"        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline\n",
						"        extra_step_kwargs = {\"generator\": generator, \"eta\": eta}\n",
						"\n",
						"        # 7. Denoising loop\n",
						"        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n",
						"        with self.progress_bar(total=num_inference_steps) as progress_bar:\n",
						"            for i, t in enumerate(timesteps):\n",
						"                # expand the latents if we are doing classifier free guidance\n",
						"                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n",
						"                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
						"\n",
						"                # predict the noise residual\n",
						"                noise_pred = self.unet(\n",
						"                    {\n",
						"                        \"sample\": latent_model_input,\n",
						"                        \"timestep\": t,\n",
						"                        \"encoder_hidden_states\": prompt_embeds,\n",
						"                    }\n",
						"                )[0]\n",
						"                noise_pred = torch.tensor(noise_pred)\n",
						"\n",
						"                # perform guidance\n",
						"                if do_classifier_free_guidance:\n",
						"                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
						"                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
						"\n",
						"                # reshape latents\n",
						"                bsz, channel, frames, width, height = latents.shape\n",
						"                latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"                noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"\n",
						"                # compute the previous noisy sample x_t -> x_t-1\n",
						"                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample\n",
						"\n",
						"                # reshape latents back\n",
						"                latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4)\n",
						"\n",
						"                # call the callback, if provided\n",
						"                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n",
						"                    progress_bar.update()\n",
						"                    if callback is not None and i % callback_steps == 0:\n",
						"                        callback(i, t, latents)\n",
						"\n",
						"        video_tensor = self.decode_latents(latents)\n",
						"\n",
						"        if output_type == \"pt\":\n",
						"            video = video_tensor\n",
						"        else:\n",
						"            video = tensor2vid(video_tensor)\n",
						"\n",
						"        if not return_dict:\n",
						"            return (video,)\n",
						"\n",
						"        return {\"frames\": video}\n",
						"\n",
						"    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt\n",
						"    def _encode_prompt(\n",
						"        self,\n",
						"        prompt,\n",
						"        num_images_per_prompt,\n",
						"        do_classifier_free_guidance,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Encodes the prompt into text encoder hidden states.\n",
						"\n",
						"        Args:\n",
						"             prompt (`str` or `List[str]`, *optional*):\n",
						"                prompt to be encoded\n",
						"            num_images_per_prompt (`int`):\n",
						"                number of images that should be generated per prompt\n",
						"            do_classifier_free_guidance (`bool`):\n",
						"                whether to use classifier free guidance or not\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"        \"\"\"\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        if prompt_embeds is None:\n",
						"            text_inputs = self.tokenizer(\n",
						"                prompt,\n",
						"                padding=\"max_length\",\n",
						"                max_length=self.tokenizer.model_max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"            text_input_ids = text_inputs.input_ids\n",
						"            untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n",
						"\n",
						"            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n",
						"                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])\n",
						"                print(\n",
						"                    \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n",
						"                    f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n",
						"                )\n",
						"\n",
						"            prompt_embeds = self.text_encoder(text_input_ids)\n",
						"            prompt_embeds = prompt_embeds[0]\n",
						"            prompt_embeds = torch.tensor(prompt_embeds)\n",
						"\n",
						"        bs_embed, seq_len, _ = prompt_embeds.shape\n",
						"        # duplicate text embeddings for each generation per prompt, using mps friendly method\n",
						"        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"        # get unconditional embeddings for classifier free guidance\n",
						"        if do_classifier_free_guidance and negative_prompt_embeds is None:\n",
						"            uncond_tokens: List[str]\n",
						"            if negative_prompt is None:\n",
						"                uncond_tokens = [\"\"] * batch_size\n",
						"            elif type(prompt) is not type(negative_prompt):\n",
						"                raise TypeError(f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\" f\" {type(prompt)}.\")\n",
						"            elif isinstance(negative_prompt, str):\n",
						"                uncond_tokens = [negative_prompt]\n",
						"            elif batch_size != len(negative_prompt):\n",
						"                raise ValueError(\n",
						"                    f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n",
						"                    f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n",
						"                    \" the batch size of `prompt`.\"\n",
						"                )\n",
						"            else:\n",
						"                uncond_tokens = negative_prompt\n",
						"\n",
						"            max_length = prompt_embeds.shape[1]\n",
						"            uncond_input = self.tokenizer(\n",
						"                uncond_tokens,\n",
						"                padding=\"max_length\",\n",
						"                max_length=max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"\n",
						"            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids)\n",
						"            negative_prompt_embeds = negative_prompt_embeds[0]\n",
						"            negative_prompt_embeds = torch.tensor(negative_prompt_embeds)\n",
						"\n",
						"        if do_classifier_free_guidance:\n",
						"            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n",
						"            seq_len = negative_prompt_embeds.shape[1]\n",
						"\n",
						"            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"            # For classifier free guidance, we need to do two forward passes.\n",
						"            # Here we concatenate the unconditional and text embeddings into a single batch\n",
						"            # to avoid doing two forward passes\n",
						"            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])\n",
						"\n",
						"        return prompt_embeds\n",
						"\n",
						"    def prepare_latents(\n",
						"        self,\n",
						"        batch_size,\n",
						"        num_channels_latents,\n",
						"        dtype,\n",
						"        generator,\n",
						"        latents=None,\n",
						"    ):\n",
						"        shape = (\n",
						"            batch_size,\n",
						"            num_channels_latents,\n",
						"            self.num_frames,\n",
						"            self.height // self.vae_scale_factor,\n",
						"            self.width // self.vae_scale_factor,\n",
						"        )\n",
						"        if isinstance(generator, list) and len(generator) != batch_size:\n",
						"            raise ValueError(\n",
						"                f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n",
						"                f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n",
						"            )\n",
						"\n",
						"        if latents is None:\n",
						"            latents = randn_tensor(shape, generator=generator, dtype=dtype)\n",
						"\n",
						"        # scale the initial noise by the standard deviation required by the scheduler\n",
						"        latents = latents * self.scheduler.init_noise_sigma\n",
						"        return latents\n",
						"\n",
						"    def check_inputs(\n",
						"        self,\n",
						"        prompt,\n",
						"        callback_steps,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds=None,\n",
						"        negative_prompt_embeds=None,\n",
						"    ):\n",
						"        if self.height % 8 != 0 or self.width % 8 != 0:\n",
						"            raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {self.height} and {self.width}.\")\n",
						"\n",
						"        if (callback_steps is None) or (callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)):\n",
						"            raise ValueError(f\"`callback_steps` has to be a positive integer but is {callback_steps} of type\" f\" {type(callback_steps)}.\")\n",
						"\n",
						"        if prompt is not None and prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to\" \" only forward one of the two.\"\n",
						"            )\n",
						"        elif prompt is None and prompt_embeds is None:\n",
						"            raise ValueError(\"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.\")\n",
						"        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):\n",
						"            raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n",
						"\n",
						"        if negative_prompt is not None and negative_prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:\"\n",
						"                f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n",
						"            )\n",
						"\n",
						"        if prompt_embeds is not None and negative_prompt_embeds is not None:\n",
						"            if prompt_embeds.shape != negative_prompt_embeds.shape:\n",
						"                raise ValueError(\n",
						"                    \"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but\"\n",
						"                    f\" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`\"\n",
						"                    f\" {negative_prompt_embeds.shape}.\"\n",
						"                )\n",
						"\n",
						"    def decode_latents(self, latents):\n",
						"        scale_factor = 0.18215\n",
						"        latents = 1 / scale_factor * latents\n",
						"\n",
						"        batch_size, channels, num_frames, height, width = latents.shape\n",
						"        latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)\n",
						"        image = self.vae_decoder(latents)[0]\n",
						"        image = torch.tensor(image)\n",
						"        video = (\n",
						"            image[None, :]\n",
						"            .reshape(\n",
						"                (\n",
						"                    batch_size,\n",
						"                    num_frames,\n",
						"                    -1,\n",
						"                )\n",
						"                + image.shape[2:]\n",
						"            )\n",
						"            .permute(0, 2, 1, 3, 4)\n",
						"        )\n",
						"        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16\n",
						"        video = video.float()\n",
						"        return video"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Inference with OpenVINO\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"device = widgets.Dropdown(\n",
						"    options=core.available_devices + [\"AUTO\"],\n",
						"    value=\"AUTO\",\n",
						"    description=\"Device:\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_unet = core.compile_model(unet_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_vae_decoder = core.compile_model(vae_decoder_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_text_encoder = core.compile_model(text_encoder_xml, device_name=device.value)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Here we replace the pipeline parts with versions converted to OpenVINO IR and compiled to specific device. Note that we use original pipeline tokenizer and scheduler."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"ov_pipe = OVTextToVideoSDPipeline(ov_vae_decoder, ov_text_encoder, tokenizer, ov_unet, scheduler)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define a prompt\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"prompt = \"A panda eating bamboo on a rock.\""
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Let's generate a video for our prompt. For full list of arguments, see `__call__` function definition of `OVTextToVideoSDPipeline` class in [Build a pipeline](#Build-a-pipeline) section."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Video generation\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"frames = ov_pipe(prompt, num_inference_steps=25)[\"frames\"]"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"images[0].save(\"output.gif\", save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"with open(\"output.gif\", \"rb\") as gif_file:\n",
						"    b64 = f\"data:image/gif;base64,{base64.b64encode(gif_file.read()).decode()}\"\n",
						"IPython.display.HTML(f'<img src=\"{b64}\" />')"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Interactive demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def generate(prompt, seed, num_inference_steps, _=gr.Progress(track_tqdm=True)):\n",
						"    generator = torch.Generator().manual_seed(seed)\n",
						"    frames = ov_pipe(\n",
						"        prompt,\n",
						"        num_inference_steps=num_inference_steps,\n",
						"        generator=generator,\n",
						"    )[\"frames\"]\n",
						"    out_file = tempfile.NamedTemporaryFile(suffix=\".gif\", delete=False)\n",
						"    images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"    images[0].save(out_file, save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"    return out_file.name\n",
						"\n",
						"\n",
						"demo = gr.Interface(\n",
						"    generate,\n",
						"    [\n",
						"        gr.Textbox(label=\"Prompt\"),\n",
						"        gr.Slider(0, 1000000, value=42, label=\"Seed\", step=1),\n",
						"        gr.Slider(10, 50, value=25, label=\"Number of inference steps\", step=1),\n",
						"    ],\n",
						"    gr.Image(label=\"Result\"),\n",
						"    examples=[\n",
						"        [\"An astronaut riding a horse.\", 0, 25],\n",
						"        [\"A panda eating bamboo on a rock.\", 0, 25],\n",
						"        [\"Spiderman is surfing.\", 0, 25],\n",
						"    ],\n",
						"    allow_flagging=\"never\",\n",
						")\n",
						"\n",
						"try:\n",
						"    demo.queue().launch(debug=True)\n",
						"except Exception:\n",
						"    demo.queue().launch(share=True, debug=True)\n",
						"# if you are launching remotely, specify server_name and server_port\n",
						"# demo.launch(server_name='your server name', server_port='server port in int')\n",
						"# Read more in the docs: https://gradio.app/docs/"
					]
				}
			]
				.map(fromJupyterCell)
			,
			[
				{
					"cell_type": "markdown",
					"source": [
						"# Video generation with ZeroScope and OpenVINO\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Install and import required packages](#Install-and-import-required-packages)\n",
						"- [Load the model](#Load-the-model)\n",
						"- [Convert the model](#Convert-the-model)\n",
						"    - [Define the conversion function](#Define-the-conversion-function)\n",
						"    - [UNet](#UNet)\n",
						"    - [VAE](#VAE)\n",
						"    - [Text encoder](#Text-encoder)\n",
						"- [Build a pipeline](#Build-a-pipeline)\n",
						"- [Inference with OpenVINO](#Inference-with-OpenVINO)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Define a prompt](#Define-a-prompt)\n",
						"    - [Video generation](#Video-generation)\n",
						"- [Interactive demo](#Interactive-demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/zeroscope-text2video/zeroscope-text2video.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The ZeroScope model is a free and open-source text-to-video model that can generate realistic and engaging videos from text descriptions. It is based on the [Modelscope](https://modelscope.cn/models/damo/text-to-video-synthesis/summary) model, but it has been improved to produce higher-quality videos with a 16:9 aspect ratio and no Shutterstock watermark. The ZeroScope model is available in two versions: ZeroScope_v2 576w, which is optimized for rapid content creation at a resolution of 576x320 pixels, and ZeroScope_v2 XL, which upscales videos to a high-definition resolution of 1024x576.\n",
						"\n",
						"The ZeroScope model is trained on a dataset of over 9,000 videos and 29,000 tagged frames. It uses a diffusion model to generate videos, which means that it starts with a random noise image and gradually adds detail to it until it matches the text description. The ZeroScope model is still under development, but it has already been used to create some impressive videos. For example, it has been used to create videos of people dancing, playing sports, and even driving cars.\n",
						"\n",
						"The ZeroScope model is a powerful tool that can be used to create various videos, from simple animations to complex scenes. It is still under development, but it has the potential to revolutionize the way we create and consume video content.\n",
						"\n",
						"Both versions of the ZeroScope model are available on Hugging Face:\n",
						" - [ZeroScope_v2 576w](https://huggingface.co/cerspense/zeroscope_v2_576w)\n",
						" - [ZeroScope_v2 XL](https://huggingface.co/cerspense/zeroscope_v2_XL)\n",
						"\n",
						"We will use the first one."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"<div class=\"alert alert-block alert-warning\">\n",
						"    This tutorial requires at least 24GB of free memory to generate a video with a frame size of 432x240 and 16 frames. Increasing either of these values will require more memory and take more time.\n",
						"</div>"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Install and import required packages\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"To work with text-to-video synthesis model, we will use Hugging Face's [Diffusers](https://github.com/huggingface/diffusers) library. It provides already pretrained model from `cerspense`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu  \"diffusers>=0.18.0\" \"torch>=2.1\" transformers \"openvino>=2023.1.0\" numpy \"gradio>=4.19\""
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import gc\n",
						"from typing import Optional, Union, List, Callable\n",
						"import base64\n",
						"import tempfile\n",
						"import warnings\n",
						"\n",
						"import diffusers\n",
						"import transformers\n",
						"import numpy as np\n",
						"import IPython\n",
						"import torch\n",
						"import PIL\n",
						"import gradio as gr\n",
						"\n",
						"import openvino as ov"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Original 576x320 inference requires a lot of RAM (>100GB), so let's run our example on a smaller frame size, keeping the same aspect ratio. Try reducing values below to reduce the memory consumption."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"WIDTH = 432  # must be divisible by 8\n",
						"HEIGHT = 240  # must be divisible by 8\n",
						"NUM_FRAMES = 16"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Load the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The model is loaded from HuggingFace using `.from_pretrained` method of `diffusers.DiffusionPipeline`."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"pipe = diffusers.DiffusionPipeline.from_pretrained(\"cerspense/zeroscope_v2_576w\")"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet = pipe.unet\n",
						"unet.eval()\n",
						"vae = pipe.vae\n",
						"vae.eval()\n",
						"text_encoder = pipe.text_encoder\n",
						"text_encoder.eval()\n",
						"tokenizer = pipe.tokenizer\n",
						"scheduler = pipe.scheduler\n",
						"vae_scale_factor = pipe.vae_scale_factor\n",
						"unet_in_channels = pipe.unet.config.in_channels\n",
						"sample_width = WIDTH // vae_scale_factor\n",
						"sample_height = HEIGHT // vae_scale_factor\n",
						"del pipe\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Convert the model\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"The architecture for generating videos from text comprises three distinct sub-networks: one for extracting text features, another for translating text features into the video latent space using a diffusion model, and a final one for mapping the video latent space to the visual space. The collective parameters of the entire model amount to approximately 1.7 billion. It's capable of processing English input. The diffusion model is built upon the Unet3D model and achieves video generation by iteratively denoising a starting point of pure Gaussian noise video."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"![](data:image/png;base64, )"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define the conversion function\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Model components are PyTorch modules, that can be converted with `ov.convert_model` function directly. We also use `ov.save_model` function to serialize the result of conversion."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"warnings.filterwarnings(\"ignore\", category=torch.jit.TracerWarning)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"\n",
						"def convert(model: torch.nn.Module, xml_path: str, **convert_kwargs) -> Path:\n",
						"    xml_path = Path(xml_path)\n",
						"    if not xml_path.exists():\n",
						"        xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
						"        with torch.no_grad():\n",
						"            converted_model = ov.convert_model(model, **convert_kwargs)\n",
						"        ov.save_model(converted_model, xml_path)\n",
						"        del converted_model\n",
						"        gc.collect()\n",
						"        torch._C._jit_clear_class_registry()\n",
						"        torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
						"        torch.jit._state._clear_class_state()\n",
						"    return xml_path"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### UNet\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text-to-video generation pipeline main component is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"unet_xml_path = convert(\n",
						"    unet,\n",
						"    \"models/unet.xml\",\n",
						"    example_input={\n",
						"        \"sample\": torch.randn(2, 4, 2, int(sample_height // 2), int(sample_width // 2)),\n",
						"        \"timestep\": torch.tensor(1),\n",
						"        \"encoder_hidden_states\": torch.randn(2, 77, 1024),\n",
						"    },\n",
						"    input=[\n",
						"        (\"sample\", (2, 4, NUM_FRAMES, sample_height, sample_width)),\n",
						"        (\"timestep\", ()),\n",
						"        (\"encoder_hidden_states\", (2, 77, 1024)),\n",
						"    ],\n",
						")\n",
						"del unet\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### VAE\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Variational autoencoder (VAE) uses UNet output to decode latents to visual representations. Our VAE model has KL loss for encoding images into latents and decoding latent representations into images. For inference, we need only decoder part."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"class VaeDecoderWrapper(torch.nn.Module):\n",
						"    def __init__(self, vae):\n",
						"        super().__init__()\n",
						"        self.vae = vae\n",
						"\n",
						"    def forward(self, z: torch.FloatTensor):\n",
						"        return self.vae.decode(z)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"vae_decoder_xml_path = convert(\n",
						"    VaeDecoderWrapper(vae),\n",
						"    \"models/vae.xml\",\n",
						"    example_input=torch.randn(2, 4, 32, 32),\n",
						"    input=((NUM_FRAMES, 4, sample_height, sample_width)),\n",
						")\n",
						"del vae\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Text encoder\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Text encoder is used to encode the input prompt to tensor. Default tensor length is 77."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"text_encoder_xml = convert(\n",
						"    text_encoder,\n",
						"    \"models/text_encoder.xml\",\n",
						"    example_input=torch.ones(1, 77, dtype=torch.int64),\n",
						"    input=((1, 77), ov.Type.i64),\n",
						")\n",
						"del text_encoder\n",
						"gc.collect();"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Build a pipeline\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def tensor2vid(video: torch.Tensor, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> List[np.ndarray]:\n",
						"    # This code is copied from https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78\n",
						"    # reshape to ncfhw\n",
						"    mean = torch.tensor(mean, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    std = torch.tensor(std, device=video.device).reshape(1, -1, 1, 1, 1)\n",
						"    # unnormalize back to [0,1]\n",
						"    video = video.mul_(std).add_(mean)\n",
						"    video.clamp_(0, 1)\n",
						"    # prepare the final outputs\n",
						"    i, c, f, h, w = video.shape\n",
						"    images = video.permute(2, 3, 0, 4, 1).reshape(f, h, i * w, c)  # 1st (frames, h, batch_size, w, c) 2nd (frames, h, batch_size * w, c)\n",
						"    images = images.unbind(dim=0)  # prepare a list of indvidual (consecutive frames)\n",
						"    images = [(image.cpu().numpy() * 255).astype(\"uint8\") for image in images]  # f h w c\n",
						"    return images"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"try:\n",
						"    from diffusers.utils import randn_tensor\n",
						"except ImportError:\n",
						"    from diffusers.utils.torch_utils import randn_tensor\n",
						"\n",
						"\n",
						"class OVTextToVideoSDPipeline(diffusers.DiffusionPipeline):\n",
						"    def __init__(\n",
						"        self,\n",
						"        vae_decoder: ov.CompiledModel,\n",
						"        text_encoder: ov.CompiledModel,\n",
						"        tokenizer: transformers.CLIPTokenizer,\n",
						"        unet: ov.CompiledModel,\n",
						"        scheduler: diffusers.schedulers.DDIMScheduler,\n",
						"    ):\n",
						"        super().__init__()\n",
						"\n",
						"        self.vae_decoder = vae_decoder\n",
						"        self.text_encoder = text_encoder\n",
						"        self.tokenizer = tokenizer\n",
						"        self.unet = unet\n",
						"        self.scheduler = scheduler\n",
						"        self.vae_scale_factor = vae_scale_factor\n",
						"        self.unet_in_channels = unet_in_channels\n",
						"        self.width = WIDTH\n",
						"        self.height = HEIGHT\n",
						"        self.num_frames = NUM_FRAMES\n",
						"\n",
						"    def __call__(\n",
						"        self,\n",
						"        prompt: Union[str, List[str]] = None,\n",
						"        num_inference_steps: int = 50,\n",
						"        guidance_scale: float = 9.0,\n",
						"        negative_prompt: Optional[Union[str, List[str]]] = None,\n",
						"        eta: float = 0.0,\n",
						"        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n",
						"        latents: Optional[torch.FloatTensor] = None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        output_type: Optional[str] = \"np\",\n",
						"        return_dict: bool = True,\n",
						"        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n",
						"        callback_steps: int = 1,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Function invoked when calling the pipeline for generation.\n",
						"\n",
						"        Args:\n",
						"            prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`.\n",
						"                instead.\n",
						"            num_inference_steps (`int`, *optional*, defaults to 50):\n",
						"                The number of denoising steps. More denoising steps usually lead to a higher quality videos at the\n",
						"                expense of slower inference.\n",
						"            guidance_scale (`float`, *optional*, defaults to 7.5):\n",
						"                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).\n",
						"                `guidance_scale` is defined as `w` of equation 2. of [Imagen\n",
						"                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n",
						"                1`. Higher guidance scale encourages to generate videos that are closely linked to the text `prompt`,\n",
						"                usually at the expense of lower video quality.\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the video generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            eta (`float`, *optional*, defaults to 0.0):\n",
						"                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n",
						"                [`schedulers.DDIMScheduler`], will be ignored for others.\n",
						"            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n",
						"                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n",
						"                to make generation deterministic.\n",
						"            latents (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video\n",
						"                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n",
						"                tensor will ge generated by sampling using the supplied random `generator`. Latents should be of shape\n",
						"                `(batch_size, num_channel, num_frames, height, width)`.\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"            output_type (`str`, *optional*, defaults to `\"np\"`):\n",
						"                The output format of the generate video. Choose between `torch.FloatTensor` or `np.array`.\n",
						"            return_dict (`bool`, *optional*, defaults to `True`):\n",
						"                Whether or not to return a [`~pipelines.stable_diffusion.TextToVideoSDPipelineOutput`] instead of a\n",
						"                plain tuple.\n",
						"            callback (`Callable`, *optional*):\n",
						"                A function that will be called every `callback_steps` steps during inference. The function will be\n",
						"                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n",
						"            callback_steps (`int`, *optional*, defaults to 1):\n",
						"                The frequency at which the `callback` function will be called. If not specified, the callback will be\n",
						"                called at every step.\n",
						"\n",
						"        Returns:\n",
						"            `List[np.ndarray]`: generated video frames\n",
						"        \"\"\"\n",
						"\n",
						"        num_images_per_prompt = 1\n",
						"\n",
						"        # 1. Check inputs. Raise error if not correct\n",
						"        self.check_inputs(\n",
						"            prompt,\n",
						"            callback_steps,\n",
						"            negative_prompt,\n",
						"            prompt_embeds,\n",
						"            negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 2. Define call parameters\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n",
						"        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n",
						"        # corresponds to doing no classifier free guidance.\n",
						"        do_classifier_free_guidance = guidance_scale > 1.0\n",
						"\n",
						"        # 3. Encode input prompt\n",
						"        prompt_embeds = self._encode_prompt(\n",
						"            prompt,\n",
						"            num_images_per_prompt,\n",
						"            do_classifier_free_guidance,\n",
						"            negative_prompt,\n",
						"            prompt_embeds=prompt_embeds,\n",
						"            negative_prompt_embeds=negative_prompt_embeds,\n",
						"        )\n",
						"\n",
						"        # 4. Prepare timesteps\n",
						"        self.scheduler.set_timesteps(num_inference_steps)\n",
						"        timesteps = self.scheduler.timesteps\n",
						"\n",
						"        # 5. Prepare latent variables\n",
						"        num_channels_latents = self.unet_in_channels\n",
						"        latents = self.prepare_latents(\n",
						"            batch_size * num_images_per_prompt,\n",
						"            num_channels_latents,\n",
						"            prompt_embeds.dtype,\n",
						"            generator,\n",
						"            latents,\n",
						"        )\n",
						"\n",
						"        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline\n",
						"        extra_step_kwargs = {\"generator\": generator, \"eta\": eta}\n",
						"\n",
						"        # 7. Denoising loop\n",
						"        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n",
						"        with self.progress_bar(total=num_inference_steps) as progress_bar:\n",
						"            for i, t in enumerate(timesteps):\n",
						"                # expand the latents if we are doing classifier free guidance\n",
						"                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n",
						"                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
						"\n",
						"                # predict the noise residual\n",
						"                noise_pred = self.unet(\n",
						"                    {\n",
						"                        \"sample\": latent_model_input,\n",
						"                        \"timestep\": t,\n",
						"                        \"encoder_hidden_states\": prompt_embeds,\n",
						"                    }\n",
						"                )[0]\n",
						"                noise_pred = torch.tensor(noise_pred)\n",
						"\n",
						"                # perform guidance\n",
						"                if do_classifier_free_guidance:\n",
						"                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
						"                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
						"\n",
						"                # reshape latents\n",
						"                bsz, channel, frames, width, height = latents.shape\n",
						"                latents = latents.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"                noise_pred = noise_pred.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channel, width, height)\n",
						"\n",
						"                # compute the previous noisy sample x_t -> x_t-1\n",
						"                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample\n",
						"\n",
						"                # reshape latents back\n",
						"                latents = latents[None, :].reshape(bsz, frames, channel, width, height).permute(0, 2, 1, 3, 4)\n",
						"\n",
						"                # call the callback, if provided\n",
						"                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):\n",
						"                    progress_bar.update()\n",
						"                    if callback is not None and i % callback_steps == 0:\n",
						"                        callback(i, t, latents)\n",
						"\n",
						"        video_tensor = self.decode_latents(latents)\n",
						"\n",
						"        if output_type == \"pt\":\n",
						"            video = video_tensor\n",
						"        else:\n",
						"            video = tensor2vid(video_tensor)\n",
						"\n",
						"        if not return_dict:\n",
						"            return (video,)\n",
						"\n",
						"        return {\"frames\": video}\n",
						"\n",
						"    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt\n",
						"    def _encode_prompt(\n",
						"        self,\n",
						"        prompt,\n",
						"        num_images_per_prompt,\n",
						"        do_classifier_free_guidance,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
						"    ):\n",
						"        r\"\"\"\n",
						"        Encodes the prompt into text encoder hidden states.\n",
						"\n",
						"        Args:\n",
						"             prompt (`str` or `List[str]`, *optional*):\n",
						"                prompt to be encoded\n",
						"            num_images_per_prompt (`int`):\n",
						"                number of images that should be generated per prompt\n",
						"            do_classifier_free_guidance (`bool`):\n",
						"                whether to use classifier free guidance or not\n",
						"            negative_prompt (`str` or `List[str]`, *optional*):\n",
						"                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
						"                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
						"                less than `1`).\n",
						"            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
						"                provided, text embeddings will be generated from `prompt` input argument.\n",
						"            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
						"                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
						"                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
						"                argument.\n",
						"        \"\"\"\n",
						"        if prompt is not None and isinstance(prompt, str):\n",
						"            batch_size = 1\n",
						"        elif prompt is not None and isinstance(prompt, list):\n",
						"            batch_size = len(prompt)\n",
						"        else:\n",
						"            batch_size = prompt_embeds.shape[0]\n",
						"\n",
						"        if prompt_embeds is None:\n",
						"            text_inputs = self.tokenizer(\n",
						"                prompt,\n",
						"                padding=\"max_length\",\n",
						"                max_length=self.tokenizer.model_max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"            text_input_ids = text_inputs.input_ids\n",
						"            untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n",
						"\n",
						"            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):\n",
						"                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])\n",
						"                print(\n",
						"                    \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n",
						"                    f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n",
						"                )\n",
						"\n",
						"            prompt_embeds = self.text_encoder(text_input_ids)\n",
						"            prompt_embeds = prompt_embeds[0]\n",
						"            prompt_embeds = torch.tensor(prompt_embeds)\n",
						"\n",
						"        bs_embed, seq_len, _ = prompt_embeds.shape\n",
						"        # duplicate text embeddings for each generation per prompt, using mps friendly method\n",
						"        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"        # get unconditional embeddings for classifier free guidance\n",
						"        if do_classifier_free_guidance and negative_prompt_embeds is None:\n",
						"            uncond_tokens: List[str]\n",
						"            if negative_prompt is None:\n",
						"                uncond_tokens = [\"\"] * batch_size\n",
						"            elif type(prompt) is not type(negative_prompt):\n",
						"                raise TypeError(f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\" f\" {type(prompt)}.\")\n",
						"            elif isinstance(negative_prompt, str):\n",
						"                uncond_tokens = [negative_prompt]\n",
						"            elif batch_size != len(negative_prompt):\n",
						"                raise ValueError(\n",
						"                    f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n",
						"                    f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n",
						"                    \" the batch size of `prompt`.\"\n",
						"                )\n",
						"            else:\n",
						"                uncond_tokens = negative_prompt\n",
						"\n",
						"            max_length = prompt_embeds.shape[1]\n",
						"            uncond_input = self.tokenizer(\n",
						"                uncond_tokens,\n",
						"                padding=\"max_length\",\n",
						"                max_length=max_length,\n",
						"                truncation=True,\n",
						"                return_tensors=\"pt\",\n",
						"            )\n",
						"\n",
						"            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids)\n",
						"            negative_prompt_embeds = negative_prompt_embeds[0]\n",
						"            negative_prompt_embeds = torch.tensor(negative_prompt_embeds)\n",
						"\n",
						"        if do_classifier_free_guidance:\n",
						"            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n",
						"            seq_len = negative_prompt_embeds.shape[1]\n",
						"\n",
						"            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
						"            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)\n",
						"\n",
						"            # For classifier free guidance, we need to do two forward passes.\n",
						"            # Here we concatenate the unconditional and text embeddings into a single batch\n",
						"            # to avoid doing two forward passes\n",
						"            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])\n",
						"\n",
						"        return prompt_embeds\n",
						"\n",
						"    def prepare_latents(\n",
						"        self,\n",
						"        batch_size,\n",
						"        num_channels_latents,\n",
						"        dtype,\n",
						"        generator,\n",
						"        latents=None,\n",
						"    ):\n",
						"        shape = (\n",
						"            batch_size,\n",
						"            num_channels_latents,\n",
						"            self.num_frames,\n",
						"            self.height // self.vae_scale_factor,\n",
						"            self.width // self.vae_scale_factor,\n",
						"        )\n",
						"        if isinstance(generator, list) and len(generator) != batch_size:\n",
						"            raise ValueError(\n",
						"                f\"You have passed a list of generators of length {len(generator)}, but requested an effective batch\"\n",
						"                f\" size of {batch_size}. Make sure the batch size matches the length of the generators.\"\n",
						"            )\n",
						"\n",
						"        if latents is None:\n",
						"            latents = randn_tensor(shape, generator=generator, dtype=dtype)\n",
						"\n",
						"        # scale the initial noise by the standard deviation required by the scheduler\n",
						"        latents = latents * self.scheduler.init_noise_sigma\n",
						"        return latents\n",
						"\n",
						"    def check_inputs(\n",
						"        self,\n",
						"        prompt,\n",
						"        callback_steps,\n",
						"        negative_prompt=None,\n",
						"        prompt_embeds=None,\n",
						"        negative_prompt_embeds=None,\n",
						"    ):\n",
						"        if self.height % 8 != 0 or self.width % 8 != 0:\n",
						"            raise ValueError(f\"`height` and `width` have to be divisible by 8 but are {self.height} and {self.width}.\")\n",
						"\n",
						"        if (callback_steps is None) or (callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)):\n",
						"            raise ValueError(f\"`callback_steps` has to be a positive integer but is {callback_steps} of type\" f\" {type(callback_steps)}.\")\n",
						"\n",
						"        if prompt is not None and prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to\" \" only forward one of the two.\"\n",
						"            )\n",
						"        elif prompt is None and prompt_embeds is None:\n",
						"            raise ValueError(\"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.\")\n",
						"        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):\n",
						"            raise ValueError(f\"`prompt` has to be of type `str` or `list` but is {type(prompt)}\")\n",
						"\n",
						"        if negative_prompt is not None and negative_prompt_embeds is not None:\n",
						"            raise ValueError(\n",
						"                f\"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:\"\n",
						"                f\" {negative_prompt_embeds}. Please make sure to only forward one of the two.\"\n",
						"            )\n",
						"\n",
						"        if prompt_embeds is not None and negative_prompt_embeds is not None:\n",
						"            if prompt_embeds.shape != negative_prompt_embeds.shape:\n",
						"                raise ValueError(\n",
						"                    \"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but\"\n",
						"                    f\" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`\"\n",
						"                    f\" {negative_prompt_embeds.shape}.\"\n",
						"                )\n",
						"\n",
						"    def decode_latents(self, latents):\n",
						"        scale_factor = 0.18215\n",
						"        latents = 1 / scale_factor * latents\n",
						"\n",
						"        batch_size, channels, num_frames, height, width = latents.shape\n",
						"        latents = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width)\n",
						"        image = self.vae_decoder(latents)[0]\n",
						"        image = torch.tensor(image)\n",
						"        video = (\n",
						"            image[None, :]\n",
						"            .reshape(\n",
						"                (\n",
						"                    batch_size,\n",
						"                    num_frames,\n",
						"                    -1,\n",
						"                )\n",
						"                + image.shape[2:]\n",
						"            )\n",
						"            .permute(0, 2, 1, 3, 4)\n",
						"        )\n",
						"        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16\n",
						"        video = video.float()\n",
						"        return video"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Inference with OpenVINO\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"core = ov.Core()"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"from notebook_utils import device_widget\n",
						"\n",
						"device = device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_unet = core.compile_model(unet_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_vae_decoder = core.compile_model(vae_decoder_xml_path, device_name=device.value)"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"%%time\n",
						"ov_text_encoder = core.compile_model(text_encoder_xml, device_name=device.value)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Here we replace the pipeline parts with versions converted to OpenVINO IR and compiled to specific device. Note that we use original pipeline tokenizer and scheduler."
					]
				},
				{
					"cell_type": "code",
					"source": [
						"ov_pipe = OVTextToVideoSDPipeline(ov_vae_decoder, ov_text_encoder, tokenizer, ov_unet, scheduler)"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Define a prompt\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"prompt = \"A panda eating bamboo on a rock.\""
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"Let's generate a video for our prompt. For full list of arguments, see `__call__` function definition of `OVTextToVideoSDPipeline` class in [Build a pipeline](#Build-a-pipeline) section."
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"### Video generation\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"frames = ov_pipe(prompt, num_inference_steps=25)[\"frames\"]"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"images[0].save(\"output.gif\", save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"with open(\"output.gif\", \"rb\") as gif_file:\n",
						"    b64 = f\"data:image/gif;base64,{base64.b64encode(gif_file.read()).decode()}\"\n",
						"IPython.display.HTML(f'<img src=\"{b64}\" />')"
					]
				},
				{
					"cell_type": "markdown",
					"source": [
						"## Interactive demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "code",
					"source": [
						"def generate(prompt, seed, num_inference_steps, _=gr.Progress(track_tqdm=True)):\n",
						"    generator = torch.Generator().manual_seed(seed)\n",
						"    frames = ov_pipe(\n",
						"        prompt,\n",
						"        num_inference_steps=num_inference_steps,\n",
						"        generator=generator,\n",
						"    )[\"frames\"]\n",
						"    out_file = tempfile.NamedTemporaryFile(suffix=\".gif\", delete=False)\n",
						"    images = [PIL.Image.fromarray(frame) for frame in frames]\n",
						"    images[0].save(out_file, save_all=True, append_images=images[1:], duration=125, loop=0)\n",
						"    return out_file.name\n",
						"\n",
						"\n",
						"demo = gr.Interface(\n",
						"    generate,\n",
						"    [\n",
						"        gr.Textbox(label=\"Prompt\"),\n",
						"        gr.Slider(0, 1000000, value=42, label=\"Seed\", step=1),\n",
						"        gr.Slider(10, 50, value=25, label=\"Number of inference steps\", step=1),\n",
						"    ],\n",
						"    gr.Image(label=\"Result\"),\n",
						"    examples=[\n",
						"        [\"An astronaut riding a horse.\", 0, 25],\n",
						"        [\"A panda eating bamboo on a rock.\", 0, 25],\n",
						"        [\"Spiderman is surfing.\", 0, 25],\n",
						"    ],\n",
						"    allow_flagging=\"never\",\n",
						")\n",
						"\n",
						"try:\n",
						"    demo.queue().launch(debug=True)\n",
						"except Exception:\n",
						"    demo.queue().launch(share=True, debug=True)\n",
						"# if you are launching remotely, specify server_name and server_port\n",
						"# demo.launch(server_name='your server name', server_port='server port in int')\n",
						"# Read more in the docs: https://gradio.app/docs/"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
			{ modified: 24, original: 24 },
			{ modified: 25, original: 25 },
			{ modified: 26, original: 26 },
			{ modified: 27, original: 27 },
			{ modified: 28, original: 28 },
			{ modified: 29, original: 29 },
			{ modified: 30, original: 30 },
			{ modified: 31, original: 31 },
			{ modified: 32, original: 32 },
			{ modified: 33, original: 33 },
			{ modified: 34, original: 34 },
			{ modified: 35, original: 35 },
			{ modified: 36, original: 36 },
			{ modified: 37, original: 37 },
			{ modified: 38, original: 38 },
			{ modified: 39, original: 39 },
			{ modified: 40, original: 40 },
			{ modified: 41, original: 41 },
			{ modified: 42, original: 42 },
			{ modified: 43, original: 43 },
			{ modified: 44, original: 44 },
			{ modified: 45, original: 45 },
			{ modified: 46, original: 46 },
			{ modified: 47, original: 47 },
			{ modified: 48, original: 48 },
			{ modified: 49, original: 49 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 6, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 15, originalLength: 1, modifiedStart: 15, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 36, originalLength: 1, modifiedStart: 36, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 50, originalLength: 0, modifiedStart: 50, modifiedLength: 2 } satisfies IDiffChange,
		]);

	});
	test('Modification of multiple cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"id": "a2e7d62b-5779-4211-822c-457c77321f8b",
					"source": [
						"# Convert and Optimize YOLOv11 instance segmentation model with OpenVINO™\n",
						"\n",
						"Instance segmentation goes a step further than object detection and involves identifying individual objects in an image and segmenting them from the rest of the image. Instance segmentation as an object detection are often used as key components in computer vision systems. \n",
						"Applications that use real-time instance segmentation models include video analytics, robotics, autonomous vehicles, multi-object tracking and object counting, medical image analysis, and many others.\n",
						"\n",
						"\n",
						"This tutorial demonstrates step-by-step instructions on how to run and optimize PyTorch YOLOv11 with OpenVINO. We consider the steps required for instance segmentation scenario.  You can find more details about model on [model page](https://docs.ultralytics.com/models/yolo11/) in Ultralytics documentation.\n",
						"\n",
						"The tutorial consists of the following steps:\n",
						"- Prepare the PyTorch model.\n",
						"- Download and prepare a dataset.\n",
						"- Validate the original model.\n",
						"- Convert the PyTorch model to OpenVINO IR.\n",
						"- Validate the converted model.\n",
						"- Prepare and run optimization pipeline.\n",
						"- Compare performance of the FP32 and quantized models.\n",
						"- Compare accuracy of the FP32 and quantized models.\n",
						"- Live demo\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Get PyTorch model](#Get-PyTorch-model)\n",
						"    - [Prerequisites](#Prerequisites)\n",
						"- [Instantiate model](#Instantiate-model)\n",
						"    - [Convert model to OpenVINO IR](#Convert-model-to-OpenVINO-IR)\n",
						"    - [Verify model inference](#Verify-model-inference)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Test on single image](#Test-on-single-image)\n",
						"- [Optimize model using NNCF Post-training Quantization API](#Optimize-model-using-NNCF-Post-training-Quantization-API)\n",
						"    - [Validate Quantized model inference](#Validate-Quantized-model-inference)\n",
						"- [Compare the Original and Quantized Models](#Compare-the-Original-and-Quantized-Models)\n",
						"    - [Compare performance of the Original and Quantized Models](#Compare-performance-of-the-Original-and-Quantized-Models)\n",
						"- [Other ways to optimize model](#Other-ways-to-optimize-model)\n",
						"- [Live demo](#Live-demo)\n",
						"    - [Run Live Object Detection and Segmentation](#Run-Live-Object-Detection-and-Segmentation)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/yolov11-optimization/yolov11-instance-segmentation.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"id": "d7a12678-b12f-48d1-9735-398855733e46",
					"source": [
						"## Get PyTorch model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Generally, PyTorch models represent an instance of the [`torch.nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) class, initialized by a state dictionary with model weights.\n",
						"We will use the YOLOv11 nano model (also known as `yolo11n-seg`) pre-trained on a COCO dataset, which is available in this [repo](https://github.com/ultralytics/ultralytics). Similar steps are also applicable to other YOLOv11 models.\n",
						"Typical steps to obtain a pre-trained model:\n",
						"1. Create an instance of a model class.\n",
						"2. Load a checkpoint state dict, which contains the pre-trained model weights.\n",
						"3. Turn the model to evaluation for switching some operations to inference mode.\n",
						"\n",
						"In this case, the creators of the model provide an API that enables converting the YOLOv11 model to OpenVINO IR. Therefore, we do not need to do these steps manually."
					]
				},
				{
					"cell_type": "markdown",
					"id": "e2267760-cbfe-41c6-958d-cad9f845d5bb",
					"source": [
						"#### Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Install necessary packages."
					]
				},
				{
					"cell_type": "code",
					"id": "30d04872-6916-454c-9211-6c644b50dc04",
					"source": [
						"%pip install -q \"openvino>=2024.0.0\" \"nncf>=2.9.0\"\n",
						"%pip install -q \"torch>=2.1\" \"torchvision>=0.16\" \"ultralytics==8.3.0\" opencv-python tqdm --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "markdown",
					"id": "1bbe319c",
					"source": [
						"Import required utility functions.\n",
						"The lower cell will download the `notebook_utils` Python module from GitHub."
					]
				},
				{
					"cell_type": "code",
					"id": "a2f6cd89",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"from notebook_utils import download_file, VideoPlayer, device_widget"
					]
				},
				{
					"cell_type": "code",
					"id": "373658bd-7e64-4479-914e-f2742d330afd",
					"source": [
						"# Download a test sample\n",
						"IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
						"download_file(\n",
						"    url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
						"    filename=IMAGE_PATH.name,\n",
						"    directory=IMAGE_PATH.parent,\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "ee32fd08-650c-4751-bb41-d8afccb2495e",
					"source": [
						"## Instantiate model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"For loading the model, required to specify a path to the model checkpoint. It can be some local path or name available on models hub (in this case model checkpoint will be downloaded automatically). \n",
						"You can select model using widget bellow:"
					]
				},
				{
					"cell_type": "code",
					"id": "7c3963a5-b50b-4ffe-b710-ccd1f4a65380",
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"model_id = [\n",
						"    \"yolo11n-seg\",\n",
						"    \"yolo11s-seg\",\n",
						"    \"yolo11m-seg\",\n",
						"    \"yolo11l-seg\",\n",
						"    \"yolo11x-seg\",\n",
						"    \"yolov8n-seg\",\n",
						"    \"yolov8s-seg\",\n",
						"    \"yolov8m-seg\",\n",
						"    \"yolov8l-seg\",\n",
						"    \"yolov8x-seg\",\n",
						"]\n",
						"\n",
						"model_name = widgets.Dropdown(options=model_id, value=model_id[0], description=\"Model\")\n",
						"\n",
						"model_name"
					]
				},
				{
					"cell_type": "markdown",
					"id": "60105d86-d7e0-47a9-ad88-4fff014ba3d8",
					"source": [
						"Making prediction, the model accepts a path to input image and returns list with Results class object. Results contains boxes for object detection model and boxes and masks for segmentation model. Also it contains utilities for processing results, for example, `plot()` method for drawing.\n",
						"\n",
						"Let us consider the examples:"
					]
				},
				{
					"cell_type": "code",
					"id": "d4994e1e-a3ee-4620-bc8b-237a55c47742",
					"source": [
						"from PIL import Image\n",
						"from ultralytics import YOLO\n",
						"\n",
						"SEG_MODEL_NAME = model_name.value\n",
						"\n",
						"seg_model = YOLO(f\"{SEG_MODEL_NAME}.pt\")\n",
						"label_map = seg_model.model.names\n",
						"\n",
						"res = seg_model(IMAGE_PATH)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "e345ffcc-c4b8-44ba-8b03-f37e63a060da",
					"source": [
						"### Convert model to OpenVINO IR\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Ultralytics provides API for convenient model exporting to different formats including OpenVINO IR. `model.export` is responsible for model conversion. We need to specify the format, and additionally, we can preserve dynamic shapes in the model."
					]
				},
				{
					"cell_type": "code",
					"id": "5c3ef363-f6a1-4fe5-b7ff-eb5a9c0789f1",
					"source": [
						"# instance segmentation model\n",
						"seg_model_path = Path(f\"{SEG_MODEL_NAME}_openvino_model/{SEG_MODEL_NAME}.xml\")\n",
						"if not seg_model_path.exists():\n",
						"    seg_model.export(format=\"openvino\", dynamic=True, half=True)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "713cd45f-2b19-4a1e-bc9d-5e69bd95d896",
					"source": [
						"### Verify model inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We can reuse the base model pipeline for pre- and postprocessing just replacing the inference method where we will use the IR model for inference."
					]
				},
				{
					"cell_type": "markdown",
					"id": "f9b9c472",
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"id": "4e0652b2",
					"source": [
						"device = device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "markdown",
					"id": "cd163eab-e803-4ae8-96da-22c3bf303630",
					"source": [
						"### Test on single image\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "34a65afa-4201-428b-af11-b35e0d206e93",
					"source": [
						"import openvino as ov\n",
						"\n",
						"core = ov.Core()\n",
						"seg_ov_model = core.read_model(seg_model_path)\n",
						"\n",
						"ov_config = {}\n",
						"if device.value != \"CPU\":\n",
						"    seg_ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"if \"GPU\" in device.value or (\"AUTO\" in device.value and \"GPU\" in core.available_devices):\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"seg_compiled_model = core.compile_model(seg_ov_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"id": "389a4698-b7fc-4aa2-9833-9681f575d88b",
					"source": [
						"seg_model = YOLO(seg_model_path.parent, task=\"segment\")\n",
						"\n",
						"if seg_model.predictor is None:\n",
						"    custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"    args = {**seg_model.overrides, **custom}\n",
						"    seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"    seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"seg_model.predictor.model.ov_compiled_model = seg_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"id": "12d2522f-e939-4512-b74a-3e20baf1edc8",
					"source": [
						"res = seg_model(IMAGE_PATH)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "7e918cf6-cdce-4e90-a6d7-c435a3c08d93",
					"source": [
						"Great! The result is the same, as produced by original models."
					]
				},
				{
					"cell_type": "markdown",
					"id": "13b69e12-2e22-44c1-bfed-fdc88f6c424d",
					"source": [
						"## Optimize model using NNCF Post-training Quantization API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
						"We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize YOLOv11.\n",
						"\n",
						"The optimization process contains the following steps:\n",
						"\n",
						"1. Create a Dataset for quantization.\n",
						"2. Run `nncf.quantize` for getting an optimized model.\n",
						"3. Serialize OpenVINO IR model, using the `openvino.runtime.serialize` function."
					]
				},
				{
					"cell_type": "markdown",
					"id": "7d7eeae8-f3f4-4e95-b1a6-a23085f03ebc",
					"source": [
						"Please select below whether you would like to run quantization to improve model inference speed."
					]
				},
				{
					"cell_type": "code",
					"id": "4690b9fc-c49f-4882-b8a5-c82f842a9126",
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"int8_model_seg_path = Path(f\"{SEG_MODEL_NAME}_openvino_int8_model/{SEG_MODEL_NAME}.xml\")\n",
						"quantized_seg_model = None\n",
						"\n",
						"to_quantize = widgets.Checkbox(\n",
						"    value=True,\n",
						"    description=\"Quantization\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"to_quantize"
					]
				},
				{
					"cell_type": "markdown",
					"id": "73cfdbd4-3b65-49e0-b8ab-ceb1d0362481",
					"source": [
						"Let's load `skip magic` extension to skip quantization if `to_quantize` is not selected"
					]
				},
				{
					"cell_type": "code",
					"id": "bc6121aa-3a07-4066-a123-fe0ed0e72478",
					"source": [
						"# Fetch skip_kernel_extension module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
						")\n",
						"open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
						"\n",
						"%load_ext skip_kernel_extension"
					]
				},
				{
					"cell_type": "markdown",
					"id": "94fb2515-3645-4798-bc7c-082129ac86c0",
					"source": [
						"Reuse validation dataloader in accuracy testing for quantization. \n",
						"For that, it should be wrapped into the `nncf.Dataset` object and define a transformation function for getting only input tensors."
					]
				},
				{
					"cell_type": "code",
					"id": "fd994958-6988-4a1d-ac7f-3efbd97135cc",
					"source": [
						"# %%skip not $to_quantize.value\n",
						"\n",
						"\n",
						"import nncf\n",
						"from typing import Dict\n",
						"\n",
						"from zipfile import ZipFile\n",
						"\n",
						"from ultralytics.data.utils import DATASETS_DIR\n",
						"from ultralytics.utils import DEFAULT_CFG\n",
						"from ultralytics.cfg import get_cfg\n",
						"from ultralytics.data.converter import coco80_to_coco91_class\n",
						"from ultralytics.data.utils import check_det_dataset\n",
						"from ultralytics.utils import ops\n",
						"\n",
						"\n",
						"if not int8_model_seg_path.exists():\n",
						"    DATA_URL = \"http://images.cocodataset.org/zips/val2017.zip\"\n",
						"    LABELS_URL = \"https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels-segments.zip\"\n",
						"    CFG_URL = \"https://raw.githubusercontent.com/ultralytics/ultralytics/v8.1.0/ultralytics/cfg/datasets/coco.yaml\"\n",
						"\n",
						"    OUT_DIR = DATASETS_DIR\n",
						"\n",
						"    DATA_PATH = OUT_DIR / \"val2017.zip\"\n",
						"    LABELS_PATH = OUT_DIR / \"coco2017labels-segments.zip\"\n",
						"    CFG_PATH = OUT_DIR / \"coco.yaml\"\n",
						"\n",
						"    download_file(DATA_URL, DATA_PATH.name, DATA_PATH.parent)\n",
						"    download_file(LABELS_URL, LABELS_PATH.name, LABELS_PATH.parent)\n",
						"    download_file(CFG_URL, CFG_PATH.name, CFG_PATH.parent)\n",
						"\n",
						"    if not (OUT_DIR / \"coco/labels\").exists():\n",
						"        with ZipFile(LABELS_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR)\n",
						"        with ZipFile(DATA_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR / \"coco/images\")\n",
						"\n",
						"    args = get_cfg(cfg=DEFAULT_CFG)\n",
						"    args.data = str(CFG_PATH)\n",
						"    seg_validator = seg_model.task_map[seg_model.task][\"validator\"](args=args)\n",
						"    seg_validator.data = check_det_dataset(args.data)\n",
						"    seg_validator.stride = 32\n",
						"    seg_data_loader = seg_validator.get_dataloader(OUT_DIR / \"coco/\", 1)\n",
						"\n",
						"    seg_validator.is_coco = True\n",
						"    seg_validator.class_map = coco80_to_coco91_class()\n",
						"    seg_validator.names = label_map\n",
						"    seg_validator.metrics.names = seg_validator.names\n",
						"    seg_validator.nc = 80\n",
						"    seg_validator.nm = 32\n",
						"    seg_validator.process = ops.process_mask\n",
						"    seg_validator.plot_masks = []\n",
						"\n",
						"    def transform_fn(data_item: Dict):\n",
						"        \"\"\"\n",
						"        Quantization transform function. Extracts and preprocess input data from dataloader item for quantization.\n",
						"        Parameters:\n",
						"           data_item: Dict with data item produced by DataLoader during iteration\n",
						"        Returns:\n",
						"            input_tensor: Input data for quantization\n",
						"        \"\"\"\n",
						"        input_tensor = seg_validator.preprocess(data_item)[\"img\"].numpy()\n",
						"        return input_tensor\n",
						"\n",
						"    quantization_dataset = nncf.Dataset(seg_data_loader, transform_fn)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "91629284-e261-494d-9464-146ed7190084",
					"source": [
						"The `nncf.quantize` function provides an interface for model quantization. It requires an instance of the OpenVINO Model and quantization dataset. \n",
						"Optionally, some additional parameters for the configuration quantization process (number of samples for quantization, preset, ignored scope, etc.) can be provided. Ultralytics models contain non-ReLU activation functions, which require asymmetric quantization of activations. To achieve a better result, we will use a `mixed` quantization preset. It provides symmetric quantization of weights and asymmetric quantization of activations. For more accurate results, we should keep the operation in the postprocessing subgraph in floating point precision, using the `ignored_scope` parameter.\n",
						"\n",
						">**Note**: Model post-training quantization is time-consuming process. Be patient, it can take several minutes depending on your hardware."
					]
				},
				{
					"cell_type": "code",
					"id": "b2e8cec4-d0b3-4da0-b54c-7964e7bcbfe2",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"if not int8_model_seg_path.exists():\n",
						"    ignored_scope = nncf.IgnoredScope(  # post-processing\n",
						"        subgraphs=[\n",
						"            nncf.Subgraph(inputs=[f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat\",\n",
						"                                  f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_1\",\n",
						"                                  f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_2\",\n",
						"                                 f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_7\"],\n",
						"                          outputs=[f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_8\"])\n",
						"        ]\n",
						"    )\n",
						"\n",
						"    # Segmentation model\n",
						"    quantized_seg_model = nncf.quantize(\n",
						"        seg_ov_model,\n",
						"        quantization_dataset,\n",
						"        preset=nncf.QuantizationPreset.MIXED,\n",
						"        ignored_scope=ignored_scope\n",
						"    )\n",
						"\n",
						"    print(f\"Quantized segmentation model will be saved to {int8_model_seg_path}\")\n",
						"    ov.save_model(quantized_seg_model, str(int8_model_seg_path))"
					]
				},
				{
					"cell_type": "markdown",
					"id": "c8c92a68-3eb8-4ea5-9e35-d496f45b3cc3",
					"source": [
						"### Validate Quantized model inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"`nncf.quantize` returns the OpenVINO Model class instance, which is suitable for loading on a device for making predictions. `INT8` model input data and output result formats have no difference from the floating point model representation. Therefore, we can reuse the same `detect` function defined above for getting the `INT8` model result on the image."
					]
				},
				{
					"cell_type": "code",
					"id": "4003a3ae",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "711200ec-2f26-47cc-b62d-3802961d5711",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"if quantized_seg_model is None:\n",
						"    quantized_seg_model = core.read_model(int8_model_seg_path)\n",
						"\n",
						"ov_config = {}\n",
						"if device.value != \"CPU\":\n",
						"    quantized_seg_model.reshape({0: [1, 3, 640, 640]})\n",
						"if \"GPU\" in device.value or (\"AUTO\" in device.value and \"GPU\" in core.available_devices):\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"\n",
						"quantized_seg_compiled_model = core.compile_model(quantized_seg_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"id": "b2e3ea3b-5935-4474-9f08-f51fc688b9c0",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"\n",
						"if seg_model.predictor is None:\n",
						"    custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"    args = {**seg_model.overrides, **custom}\n",
						"    seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"    seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"seg_model.predictor.model.ov_compiled_model = quantized_seg_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"id": "bef2ed5d-5762-41a6-af2e-ee4f9fd56557",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"res = seg_model(IMAGE_PATH)\n",
						"display(Image.fromarray(res[0].plot()[:, :, ::-1]))"
					]
				},
				{
					"cell_type": "markdown",
					"id": "b289f057",
					"source": [
						"## Compare the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"cell_type": "markdown",
					"id": "59081cae-9c92-46f0-8117-ed8c6fa6ff19",
					"source": [
						"### Compare performance of the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"Finally, use the OpenVINO [Benchmark Tool](https://docs.openvino.ai/2024/learn-openvino/openvino-samples/benchmark-tool.html) to measure the inference performance of the `FP32` and `INT8` models.\n",
						"\n",
						"> **Note**: For more accurate performance, it is recommended to run `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app -m <model_path> -d CPU -shape \"<input_shape>\"` to benchmark async inference on CPU on specific input data shape for one minute. Change `CPU` to `GPU` to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
					]
				},
				{
					"cell_type": "code",
					"id": "b8677df4",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "546828d2",
					"source": [
						"if int8_model_seg_path.exists():\n",
						"    !benchmark_app -m $seg_model_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "code",
					"id": "1d43554b",
					"source": [
						"if int8_model_seg_path.exists():\n",
						"    !benchmark_app -m $int8_model_seg_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "markdown",
					"id": "184bcebd-588d-4b4d-a60a-0fa753a07ef9",
					"source": [
						"## Other ways to optimize model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The performance could be also improved by another OpenVINO method such as async inference pipeline or preprocessing API.\n",
						"\n",
						"Async Inference pipeline help to utilize the device more optimal. The key advantage of the Async API is that when a device is busy with inference, the application can perform other tasks in parallel (for example, populating inputs or scheduling other requests) rather than wait for the current inference to complete first. To understand how to perform async inference using openvino, refer to [Async API tutorial](../async-api/async-api.ipynb)\n",
						"\n",
						"Preprocessing API enables making preprocessing a part of the model reducing application code and dependency on additional image processing libraries. \n",
						"The main advantage of Preprocessing API is that preprocessing steps will be integrated into the execution graph and will be performed on a selected device (CPU/GPU etc.) rather than always being executed on CPU as part of an application. This will also improve selected device utilization. For more information, refer to the overview of [Preprocessing API tutorial](../optimize-preprocessing/optimize-preprocessing.ipynb). To see, how it could be used with YOLOV8 object detection model, please, see [Convert and Optimize YOLOv8 real-time object detection with OpenVINO tutorial](./yolov8-object-detection.ipynb)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "2e3b4862-c182-4ce4-a473-9f38d98deab8",
					"source": [
						"## Live demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The following code runs model inference on a video:"
					]
				},
				{
					"cell_type": "code",
					"id": "4c7bb92e-e301-45a9-b5ff-f7953fad298f",
					"source": [
						"import collections\n",
						"import time\n",
						"import cv2\n",
						"from IPython import display\n",
						"import numpy as np\n",
						"\n",
						"\n",
						"def run_instance_segmentation(\n",
						"    source=0,\n",
						"    flip=False,\n",
						"    use_popup=False,\n",
						"    skip_first_frames=0,\n",
						"    model=seg_model,\n",
						"    device=device.value,\n",
						"):\n",
						"    player = None\n",
						"\n",
						"    ov_config = {}\n",
						"    if device != \"CPU\":\n",
						"        model.reshape({0: [1, 3, 640, 640]})\n",
						"    if \"GPU\" in device or (\"AUTO\" in device and \"GPU\" in core.available_devices):\n",
						"        ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"    compiled_model = core.compile_model(model, device, ov_config)\n",
						"\n",
						"    if seg_model.predictor is None:\n",
						"        custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"        args = {**seg_model.overrides, **custom}\n",
						"        seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"        seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"    seg_model.predictor.model.ov_compiled_model = compiled_model\n",
						"\n",
						"    try:\n",
						"        # Create a video player to play with target fps.\n",
						"        player = VideoPlayer(source=source, flip=flip, fps=30, skip_first_frames=skip_first_frames)\n",
						"        # Start capturing.\n",
						"        player.start()\n",
						"        if use_popup:\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(winname=title, flags=cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"        while True:\n",
						"            # Grab the frame.\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"            # If the frame is larger than full HD, reduce size to improve the performance.\n",
						"            scale = 1280 / max(frame.shape)\n",
						"            if scale < 1:\n",
						"                frame = cv2.resize(\n",
						"                    src=frame,\n",
						"                    dsize=None,\n",
						"                    fx=scale,\n",
						"                    fy=scale,\n",
						"                    interpolation=cv2.INTER_AREA,\n",
						"                )\n",
						"            # Get the results.\n",
						"            input_image = np.array(frame)\n",
						"\n",
						"            start_time = time.time()\n",
						"            detections = seg_model(input_image)\n",
						"            stop_time = time.time()\n",
						"            frame = detections[0].plot()\n",
						"\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Use processing times from last 200 frames.\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            _, f_width = frame.shape[:2]\n",
						"            # Mean processing time [ms].\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"            cv2.putText(\n",
						"                img=frame,\n",
						"                text=f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                org=(20, 40),\n",
						"                fontFace=cv2.FONT_HERSHEY_COMPLEX,\n",
						"                fontScale=f_width / 1000,\n",
						"                color=(0, 0, 255),\n",
						"                thickness=1,\n",
						"                lineType=cv2.LINE_AA,\n",
						"            )\n",
						"            # Use this workaround if there is flickering.\n",
						"            if use_popup:\n",
						"                cv2.imshow(winname=title, mat=frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # Encode numpy array to jpg.\n",
						"                _, encoded_img = cv2.imencode(ext=\".jpg\", img=frame, params=[cv2.IMWRITE_JPEG_QUALITY, 100])\n",
						"                # Create an IPython image.\n",
						"                i = display.Image(data=encoded_img)\n",
						"                # Display the image in this notebook.\n",
						"                display.clear_output(wait=True)\n",
						"                display.display(i)\n",
						"    # ctrl-c\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    # any different error\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        if player is not None:\n",
						"            # Stop capturing.\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()"
					]
				},
				{
					"cell_type": "markdown",
					"id": "cc5b4ba6-f478-4417-b09d-93fee5adca41",
					"source": [
						"### Run Live Object Detection and Segmentation\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Use a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						">**NOTE**: To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a remote server (for example, in Binder or Google Colab service), the webcam will not work. By default, the lower cell will run model inference on a video file. If you want to try live inference on your webcam set `WEBCAM_INFERENCE = True`"
					]
				},
				{
					"cell_type": "code",
					"id": "90708017",
					"source": [
						"WEBCAM_INFERENCE = False\n",
						"\n",
						"if WEBCAM_INFERENCE:\n",
						"    VIDEO_SOURCE = 0  # Webcam\n",
						"else:\n",
						"    VIDEO_SOURCE = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\""
					]
				},
				{
					"cell_type": "code",
					"id": "e8f1239c",
					"source": [
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "440f2d0e",
					"source": [
						"run_instance_segmentation(\n",
						"    source=VIDEO_SOURCE,\n",
						"    flip=True,\n",
						"    use_popup=False,\n",
						"    model=seg_ov_model,\n",
						"    device=device.value,\n",
						")"
					]
				}
			]
				.map(fromJupyterCell)
			,
			[
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "a2e7d62b-5779-4211-822c-457c77321f8b",
					"metadata": {},
					"source": [
						"# Convert and Optimize YOLOv11 instance segmentation model with OpenVINO™\n",
						"\n",
						"Instance segmentation goes a step further than object detection and involves identifying individual objects in an image and segmenting them from the rest of the image. Instance segmentation as an object detection are often used as key components in computer vision systems. \n",
						"Applications that use real-time instance segmentation models include video analytics, robotics, autonomous vehicles, multi-object tracking and object counting, medical image analysis, and many others.\n",
						"\n",
						"\n",
						"This tutorial demonstrates step-by-step instructions on how to run and optimize PyTorch YOLOv11 with OpenVINO. We consider the steps required for instance segmentation scenario.  You can find more details about model on [model page](https://docs.ultralytics.com/models/yolo11/) in Ultralytics documentation.\n",
						"\n",
						"The tutorial consists of the following steps:\n",
						"- Prepare the PyTorch model.\n",
						"- Download and prepare a dataset.\n",
						"- Validate the original model.\n",
						"- Convert the PyTorch model to OpenVINO IR.\n",
						"- Validate the converted model.\n",
						"- Prepare and run optimization pipeline.\n",
						"- Compare performance of the FP32 and quantized models.\n",
						"- Compare accuracy of the FP32 and quantized models.\n",
						"- Live demo\n",
						"\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Get PyTorch model](#Get-PyTorch-model)\n",
						"    - [Prerequisites](#Prerequisites)\n",
						"- [Instantiate model](#Instantiate-model)\n",
						"    - [Convert model to OpenVINO IR](#Convert-model-to-OpenVINO-IR)\n",
						"    - [Verify model inference](#Verify-model-inference)\n",
						"    - [Select inference device](#Select-inference-device)\n",
						"    - [Test on single image](#Test-on-single-image)\n",
						"- [Optimize model using NNCF Post-training Quantization API](#Optimize-model-using-NNCF-Post-training-Quantization-API)\n",
						"    - [Validate Quantized model inference](#Validate-Quantized-model-inference)\n",
						"- [Compare the Original and Quantized Models](#Compare-the-Original-and-Quantized-Models)\n",
						"    - [Compare performance of the Original and Quantized Models](#Compare-performance-of-the-Original-and-Quantized-Models)\n",
						"- [Other ways to optimize model](#Other-ways-to-optimize-model)\n",
						"- [Live demo](#Live-demo)\n",
						"    - [Run Live Object Detection and Segmentation](#Run-Live-Object-Detection-and-Segmentation)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/yolov11-optimization/yolov11-instance-segmentation.ipynb\" />\n"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "d7a12678-b12f-48d1-9735-398855733e46",
					"metadata": {},
					"source": [
						"## Get PyTorch model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Generally, PyTorch models represent an instance of the [`torch.nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) class, initialized by a state dictionary with model weights.\n",
						"We will use the YOLOv11 nano model (also known as `yolo11n-seg`) pre-trained on a COCO dataset, which is available in this [repo](https://github.com/ultralytics/ultralytics). Similar steps are also applicable to other YOLOv11 models.\n",
						"Typical steps to obtain a pre-trained model:\n",
						"1. Create an instance of a model class.\n",
						"2. Load a checkpoint state dict, which contains the pre-trained model weights.\n",
						"3. Turn the model to evaluation for switching some operations to inference mode.\n",
						"\n",
						"In this case, the creators of the model provide an API that enables converting the YOLOv11 model to OpenVINO IR. Therefore, we do not need to do these steps manually."
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "e2267760-cbfe-41c6-958d-cad9f845d5bb",
					"metadata": {},
					"source": [
						"#### Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Install necessary packages."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "30d04872-6916-454c-9211-6c644b50dc04",
					"metadata": {},
					"outputs": [],
					"source": [
						"%pip install -q \"openvino>=2024.0.0\" \"nncf>=2.9.0\"\n",
						"%pip install -q \"torch>=2.1\" \"torchvision>=0.16\" \"ultralytics==8.3.0\" opencv-python tqdm --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "1bbe319c",
					"metadata": {},
					"source": [
						"Import required utility functions.\n",
						"The lower cell will download the `notebook_utils` Python module from GitHub."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "a2f6cd89",
					"metadata": {},
					"outputs": [],
					"source": [
						"from pathlib import Path\n",
						"import requests\n",
						"\n",
						"if not Path(\"notebook_utils.py\").exists():\n",
						"    r = requests.get(\n",
						"        url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						"    )\n",
						"\n",
						"    open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"\n",
						"from notebook_utils import download_file, VideoPlayer, device_widget"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "373658bd-7e64-4479-914e-f2742d330afd",
					"metadata": {},
					"outputs": [],
					"source": [
						"# Download a test sample\n",
						"IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
						"\n",
						"if not IMAGE_PATH.exists():\n",
						"    download_file(\n",
						"        url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
						"        filename=IMAGE_PATH.name,\n",
						"        directory=IMAGE_PATH.parent,\n",
						"    )"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "ee32fd08-650c-4751-bb41-d8afccb2495e",
					"metadata": {},
					"source": [
						"## Instantiate model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"For loading the model, required to specify a path to the model checkpoint. It can be some local path or name available on models hub (in this case model checkpoint will be downloaded automatically). \n",
						"You can select model using widget bellow:"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "7c3963a5-b50b-4ffe-b710-ccd1f4a65380",
					"metadata": {},
					"outputs": [],
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"model_id = [\n",
						"    \"yolo11n-seg\",\n",
						"    \"yolo11s-seg\",\n",
						"    \"yolo11m-seg\",\n",
						"    \"yolo11l-seg\",\n",
						"    \"yolo11x-seg\",\n",
						"    \"yolov8n-seg\",\n",
						"    \"yolov8s-seg\",\n",
						"    \"yolov8m-seg\",\n",
						"    \"yolov8l-seg\",\n",
						"    \"yolov8x-seg\",\n",
						"]\n",
						"\n",
						"model_name = widgets.Dropdown(options=model_id, value=model_id[0], description=\"Model\")\n",
						"\n",
						"model_name"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "60105d86-d7e0-47a9-ad88-4fff014ba3d8",
					"metadata": {},
					"source": [
						"Making prediction, the model accepts a path to input image and returns list with Results class object. Results contains boxes for object detection model and boxes and masks for segmentation model. Also it contains utilities for processing results, for example, `plot()` method for drawing.\n",
						"\n",
						"Let us consider the examples:"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "d4994e1e-a3ee-4620-bc8b-237a55c47742",
					"metadata": {},
					"outputs": [],
					"source": [
						"from PIL import Image\n",
						"from ultralytics import YOLO\n",
						"\n",
						"SEG_MODEL_NAME = model_name.value\n",
						"\n",
						"seg_model = YOLO(f\"{SEG_MODEL_NAME}.pt\")\n",
						"label_map = seg_model.model.names\n",
						"\n",
						"res = seg_model(IMAGE_PATH)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "e345ffcc-c4b8-44ba-8b03-f37e63a060da",
					"metadata": {},
					"source": [
						"### Convert model to OpenVINO IR\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Ultralytics provides API for convenient model exporting to different formats including OpenVINO IR. `model.export` is responsible for model conversion. We need to specify the format, and additionally, we can preserve dynamic shapes in the model."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "5c3ef363-f6a1-4fe5-b7ff-eb5a9c0789f1",
					"metadata": {},
					"outputs": [],
					"source": [
						"# instance segmentation model\n",
						"seg_model_path = Path(f\"{SEG_MODEL_NAME}_openvino_model/{SEG_MODEL_NAME}.xml\")\n",
						"if not seg_model_path.exists():\n",
						"    seg_model.export(format=\"openvino\", dynamic=True, half=True)"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "713cd45f-2b19-4a1e-bc9d-5e69bd95d896",
					"metadata": {},
					"source": [
						"### Verify model inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"We can reuse the base model pipeline for pre- and postprocessing just replacing the inference method where we will use the IR model for inference."
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "f9b9c472",
					"metadata": {},
					"source": [
						"### Select inference device\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Select device from dropdown list for running inference using OpenVINO"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "4e0652b2",
					"metadata": {},
					"outputs": [],
					"source": [
						"device = device_widget()\n",
						"\n",
						"device"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "cd163eab-e803-4ae8-96da-22c3bf303630",
					"metadata": {},
					"source": [
						"### Test on single image\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "34a65afa-4201-428b-af11-b35e0d206e93",
					"metadata": {},
					"outputs": [],
					"source": [
						"import openvino as ov\n",
						"\n",
						"core = ov.Core()\n",
						"seg_ov_model = core.read_model(seg_model_path)\n",
						"\n",
						"ov_config = {}\n",
						"if device.value != \"CPU\":\n",
						"    seg_ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"if \"GPU\" in device.value or (\"AUTO\" in device.value and \"GPU\" in core.available_devices):\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"seg_compiled_model = core.compile_model(seg_ov_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "389a4698-b7fc-4aa2-9833-9681f575d88b",
					"metadata": {},
					"outputs": [],
					"source": [
						"seg_model = YOLO(seg_model_path.parent, task=\"segment\")\n",
						"\n",
						"if seg_model.predictor is None:\n",
						"    custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"    args = {**seg_model.overrides, **custom}\n",
						"    seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"    seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"seg_model.predictor.model.ov_compiled_model = seg_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "12d2522f-e939-4512-b74a-3e20baf1edc8",
					"metadata": {},
					"outputs": [],
					"source": [
						"res = seg_model(IMAGE_PATH)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "7e918cf6-cdce-4e90-a6d7-c435a3c08d93",
					"metadata": {},
					"source": [
						"Great! The result is the same, as produced by original models."
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "13b69e12-2e22-44c1-bfed-fdc88f6c424d",
					"metadata": {},
					"source": [
						"## Optimize model using NNCF Post-training Quantization API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
						"We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize YOLOv11.\n",
						"\n",
						"The optimization process contains the following steps:\n",
						"\n",
						"1. Create a Dataset for quantization.\n",
						"2. Run `nncf.quantize` for getting an optimized model.\n",
						"3. Serialize OpenVINO IR model, using the `openvino.runtime.serialize` function."
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "7d7eeae8-f3f4-4e95-b1a6-a23085f03ebc",
					"metadata": {},
					"source": [
						"Please select below whether you would like to run quantization to improve model inference speed."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "4690b9fc-c49f-4882-b8a5-c82f842a9126",
					"metadata": {},
					"outputs": [],
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"int8_model_seg_path = Path(f\"{SEG_MODEL_NAME}_openvino_int8_model/{SEG_MODEL_NAME}.xml\")\n",
						"quantized_seg_model = None\n",
						"\n",
						"to_quantize = widgets.Checkbox(\n",
						"    value=True,\n",
						"    description=\"Quantization\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"to_quantize"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "73cfdbd4-3b65-49e0-b8ab-ceb1d0362481",
					"metadata": {},
					"source": [
						"Let's load `skip magic` extension to skip quantization if `to_quantize` is not selected"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "bc6121aa-3a07-4066-a123-fe0ed0e72478",
					"metadata": {},
					"outputs": [],
					"source": [
						"# Fetch skip_kernel_extension module\n",
						"import requests\n",
						"\n",
						"\n",
						"if not Path(\"skip_kernel_extension.py\").exists():\n",
						"    r = requests.get(\n",
						"        url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
						"    )\n",
						"    open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
						"\n",
						"%load_ext skip_kernel_extension"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "94fb2515-3645-4798-bc7c-082129ac86c0",
					"metadata": {},
					"source": [
						"Reuse validation dataloader in accuracy testing for quantization. \n",
						"For that, it should be wrapped into the `nncf.Dataset` object and define a transformation function for getting only input tensors."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "fd994958-6988-4a1d-ac7f-3efbd97135cc",
					"metadata": {},
					"outputs": [],
					"source": [
						"# %%skip not $to_quantize.value\n",
						"\n",
						"\n",
						"import nncf\n",
						"from typing import Dict\n",
						"\n",
						"from zipfile import ZipFile\n",
						"\n",
						"from ultralytics.data.utils import DATASETS_DIR\n",
						"from ultralytics.utils import DEFAULT_CFG\n",
						"from ultralytics.cfg import get_cfg\n",
						"from ultralytics.data.converter import coco80_to_coco91_class\n",
						"from ultralytics.data.utils import check_det_dataset\n",
						"from ultralytics.utils import ops\n",
						"\n",
						"\n",
						"if not int8_model_seg_path.exists():\n",
						"    DATA_URL = \"http://images.cocodataset.org/zips/val2017.zip\"\n",
						"    LABELS_URL = \"https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels-segments.zip\"\n",
						"    CFG_URL = \"https://raw.githubusercontent.com/ultralytics/ultralytics/v8.1.0/ultralytics/cfg/datasets/coco.yaml\"\n",
						"\n",
						"    OUT_DIR = DATASETS_DIR\n",
						"\n",
						"    DATA_PATH = OUT_DIR / \"val2017.zip\"\n",
						"    LABELS_PATH = OUT_DIR / \"coco2017labels-segments.zip\"\n",
						"    CFG_PATH = OUT_DIR / \"coco.yaml\"\n",
						"\n",
						"    download_file(DATA_URL, DATA_PATH.name, DATA_PATH.parent)\n",
						"    download_file(LABELS_URL, LABELS_PATH.name, LABELS_PATH.parent)\n",
						"    download_file(CFG_URL, CFG_PATH.name, CFG_PATH.parent)\n",
						"\n",
						"    if not (OUT_DIR / \"coco/labels\").exists():\n",
						"        with ZipFile(LABELS_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR)\n",
						"        with ZipFile(DATA_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR / \"coco/images\")\n",
						"\n",
						"    args = get_cfg(cfg=DEFAULT_CFG)\n",
						"    args.data = str(CFG_PATH)\n",
						"    seg_validator = seg_model.task_map[seg_model.task][\"validator\"](args=args)\n",
						"    seg_validator.data = check_det_dataset(args.data)\n",
						"    seg_validator.stride = 32\n",
						"    seg_data_loader = seg_validator.get_dataloader(OUT_DIR / \"coco/\", 1)\n",
						"\n",
						"    seg_validator.is_coco = True\n",
						"    seg_validator.class_map = coco80_to_coco91_class()\n",
						"    seg_validator.names = label_map\n",
						"    seg_validator.metrics.names = seg_validator.names\n",
						"    seg_validator.nc = 80\n",
						"    seg_validator.nm = 32\n",
						"    seg_validator.process = ops.process_mask\n",
						"    seg_validator.plot_masks = []\n",
						"\n",
						"    def transform_fn(data_item: Dict):\n",
						"        \"\"\"\n",
						"        Quantization transform function. Extracts and preprocess input data from dataloader item for quantization.\n",
						"        Parameters:\n",
						"           data_item: Dict with data item produced by DataLoader during iteration\n",
						"        Returns:\n",
						"            input_tensor: Input data for quantization\n",
						"        \"\"\"\n",
						"        input_tensor = seg_validator.preprocess(data_item)[\"img\"].numpy()\n",
						"        return input_tensor\n",
						"\n",
						"    quantization_dataset = nncf.Dataset(seg_data_loader, transform_fn)"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "91629284-e261-494d-9464-146ed7190084",
					"metadata": {},
					"source": [
						"The `nncf.quantize` function provides an interface for model quantization. It requires an instance of the OpenVINO Model and quantization dataset. \n",
						"Optionally, some additional parameters for the configuration quantization process (number of samples for quantization, preset, ignored scope, etc.) can be provided. Ultralytics models contain non-ReLU activation functions, which require asymmetric quantization of activations. To achieve a better result, we will use a `mixed` quantization preset. It provides symmetric quantization of weights and asymmetric quantization of activations. For more accurate results, we should keep the operation in the postprocessing subgraph in floating point precision, using the `ignored_scope` parameter.\n",
						"\n",
						">**Note**: Model post-training quantization is time-consuming process. Be patient, it can take several minutes depending on your hardware."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "b2e8cec4-d0b3-4da0-b54c-7964e7bcbfe2",
					"metadata": {
						"test_replace": {
							"ignored_scope=ignored_scope\n": "ignored_scope=ignored_scope, subset_size=10\n"
						}
					},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"if not int8_model_seg_path.exists():\n",
						"    ignored_scope = nncf.IgnoredScope(  # post-processing\n",
						"        subgraphs=[\n",
						"            nncf.Subgraph(inputs=[f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat\",\n",
						"                                  f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_1\",\n",
						"                                  f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_2\",\n",
						"                                 f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_7\"],\n",
						"                          outputs=[f\"__module.model.{22 if 'v8' in SEG_MODEL_NAME else 23}/aten::cat/Concat_8\"])\n",
						"        ]\n",
						"    )\n",
						"\n",
						"    # Segmentation model\n",
						"    quantized_seg_model = nncf.quantize(\n",
						"        seg_ov_model,\n",
						"        quantization_dataset,\n",
						"        preset=nncf.QuantizationPreset.MIXED,\n",
						"        ignored_scope=ignored_scope\n",
						"    )\n",
						"\n",
						"    print(f\"Quantized segmentation model will be saved to {int8_model_seg_path}\")\n",
						"    ov.save_model(quantized_seg_model, str(int8_model_seg_path))"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "c8c92a68-3eb8-4ea5-9e35-d496f45b3cc3",
					"metadata": {},
					"source": [
						"### Validate Quantized model inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"`nncf.quantize` returns the OpenVINO Model class instance, which is suitable for loading on a device for making predictions. `INT8` model input data and output result formats have no difference from the floating point model representation. Therefore, we can reuse the same `detect` function defined above for getting the `INT8` model result on the image."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "4003a3ae",
					"metadata": {},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "711200ec-2f26-47cc-b62d-3802961d5711",
					"metadata": {},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"if quantized_seg_model is None:\n",
						"    quantized_seg_model = core.read_model(int8_model_seg_path)\n",
						"\n",
						"ov_config = {}\n",
						"if device.value != \"CPU\":\n",
						"    quantized_seg_model.reshape({0: [1, 3, 640, 640]})\n",
						"if \"GPU\" in device.value or (\"AUTO\" in device.value and \"GPU\" in core.available_devices):\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"\n",
						"quantized_seg_compiled_model = core.compile_model(quantized_seg_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "b2e3ea3b-5935-4474-9f08-f51fc688b9c0",
					"metadata": {},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"\n",
						"if seg_model.predictor is None:\n",
						"    custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"    args = {**seg_model.overrides, **custom}\n",
						"    seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"    seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"seg_model.predictor.model.ov_compiled_model = quantized_seg_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "bef2ed5d-5762-41a6-af2e-ee4f9fd56557",
					"metadata": {},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"res = seg_model(IMAGE_PATH)\n",
						"display(Image.fromarray(res[0].plot()[:, :, ::-1]))"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "b289f057",
					"metadata": {},
					"source": [
						"## Compare the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)\n"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "59081cae-9c92-46f0-8117-ed8c6fa6ff19",
					"metadata": {},
					"source": [
						"### Compare performance of the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"Finally, use the OpenVINO [Benchmark Tool](https://docs.openvino.ai/2024/learn-openvino/openvino-samples/benchmark-tool.html) to measure the inference performance of the `FP32` and `INT8` models.\n",
						"\n",
						"> **Note**: For more accurate performance, it is recommended to run `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app -m <model_path> -d CPU -shape \"<input_shape>\"` to benchmark async inference on CPU on specific input data shape for one minute. Change `CPU` to `GPU` to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "b8677df4",
					"metadata": {},
					"outputs": [],
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "546828d2",
					"metadata": {},
					"outputs": [],
					"source": [
						"if int8_model_seg_path.exists():\n",
						"    !benchmark_app -m $seg_model_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "1d43554b",
					"metadata": {},
					"outputs": [],
					"source": [
						"if int8_model_seg_path.exists():\n",
						"    !benchmark_app -m $int8_model_seg_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "184bcebd-588d-4b4d-a60a-0fa753a07ef9",
					"metadata": {},
					"source": [
						"## Other ways to optimize model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The performance could be also improved by another OpenVINO method such as async inference pipeline or preprocessing API.\n",
						"\n",
						"Async Inference pipeline help to utilize the device more optimal. The key advantage of the Async API is that when a device is busy with inference, the application can perform other tasks in parallel (for example, populating inputs or scheduling other requests) rather than wait for the current inference to complete first. To understand how to perform async inference using openvino, refer to [Async API tutorial](../async-api/async-api.ipynb)\n",
						"\n",
						"Preprocessing API enables making preprocessing a part of the model reducing application code and dependency on additional image processing libraries. \n",
						"The main advantage of Preprocessing API is that preprocessing steps will be integrated into the execution graph and will be performed on a selected device (CPU/GPU etc.) rather than always being executed on CPU as part of an application. This will also improve selected device utilization. For more information, refer to the overview of [Preprocessing API tutorial](../optimize-preprocessing/optimize-preprocessing.ipynb). To see, how it could be used with YOLOV8 object detection model, please, see [Convert and Optimize YOLOv8 real-time object detection with OpenVINO tutorial](./yolov8-object-detection.ipynb)"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "2e3b4862-c182-4ce4-a473-9f38d98deab8",
					"metadata": {
						"tags": []
					},
					"source": [
						"## Live demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"The following code runs model inference on a video:"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "4c7bb92e-e301-45a9-b5ff-f7953fad298f",
					"metadata": {},
					"outputs": [],
					"source": [
						"import collections\n",
						"import time\n",
						"import cv2\n",
						"from IPython import display\n",
						"import numpy as np\n",
						"\n",
						"\n",
						"def run_instance_segmentation(\n",
						"    source=0,\n",
						"    flip=False,\n",
						"    use_popup=False,\n",
						"    skip_first_frames=0,\n",
						"    model=seg_model,\n",
						"    device=device.value,\n",
						"    video_width=1280,\n",
						"):\n",
						"    player = None\n",
						"\n",
						"    ov_config = {}\n",
						"    if device != \"CPU\":\n",
						"        model.reshape({0: [1, 3, 640, 640]})\n",
						"    if \"GPU\" in device or (\"AUTO\" in device and \"GPU\" in core.available_devices):\n",
						"        ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"    compiled_model = core.compile_model(model, device, ov_config)\n",
						"\n",
						"    if seg_model.predictor is None:\n",
						"        custom = {\"conf\": 0.25, \"batch\": 1, \"save\": False, \"mode\": \"predict\"}  # method defaults\n",
						"        args = {**seg_model.overrides, **custom}\n",
						"        seg_model.predictor = seg_model._smart_load(\"predictor\")(overrides=args, _callbacks=seg_model.callbacks)\n",
						"        seg_model.predictor.setup_model(model=seg_model.model)\n",
						"\n",
						"    seg_model.predictor.model.ov_compiled_model = compiled_model\n",
						"\n",
						"    try:\n",
						"        # Create a video player to play with target fps.\n",
						"        player = VideoPlayer(source=source, flip=flip, fps=30, skip_first_frames=skip_first_frames, width=video_width)\n",
						"        # Start capturing.\n",
						"        player.start()\n",
						"        if use_popup:\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(winname=title, flags=cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"        while True:\n",
						"            # Grab the frame.\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"            # If the frame is larger than video_width, reduce size to improve the performance.\n",
						"            # If more, increase size for better demo expirience.\n",
						"            scale = video_width / max(frame.shape)\n",
						"            frame = cv2.resize(\n",
						"                src=frame,\n",
						"                dsize=None,\n",
						"                fx=scale,\n",
						"                fy=scale,\n",
						"                interpolation=cv2.INTER_AREA,\n",
						"            )\n",
						"            # Get the results.\n",
						"            input_image = np.array(frame)\n",
						"\n",
						"            start_time = time.time()\n",
						"            detections = seg_model(input_image)\n",
						"            stop_time = time.time()\n",
						"            frame = detections[0].plot()\n",
						"\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Use processing times from last 200 frames.\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            _, f_width = frame.shape[:2]\n",
						"            # Mean processing time [ms].\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"            cv2.putText(\n",
						"                img=frame,\n",
						"                text=f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                org=(20, 40),\n",
						"                fontFace=cv2.FONT_HERSHEY_COMPLEX,\n",
						"                fontScale=f_width / 1000,\n",
						"                color=(0, 0, 255),\n",
						"                thickness=1,\n",
						"                lineType=cv2.LINE_AA,\n",
						"            )\n",
						"            # Use this workaround if there is flickering.\n",
						"            if use_popup:\n",
						"                cv2.imshow(winname=title, mat=frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # Encode numpy array to jpg.\n",
						"                _, encoded_img = cv2.imencode(ext=\".jpg\", img=frame, params=[cv2.IMWRITE_JPEG_QUALITY, 100])\n",
						"                # Create an IPython image.\n",
						"                i = display.Image(data=encoded_img)\n",
						"                # Display the image in this notebook.\n",
						"                display.clear_output(wait=True)\n",
						"                display.display(i)\n",
						"    # ctrl-c\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    # any different error\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        if player is not None:\n",
						"            # Stop capturing.\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()"
					]
				},
				{
					"attachments": {},
					"cell_type": "markdown",
					"id": "cc5b4ba6-f478-4417-b09d-93fee5adca41",
					"metadata": {},
					"source": [
						"### Run Live Object Detection and Segmentation\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"Use a webcam as the video input. By default, the primary webcam is set with `source=0`. If you have multiple webcams, each one will be assigned a consecutive number starting at 0. Set `flip=True` when using a front-facing camera. Some web browsers, especially Mozilla Firefox, may cause flickering. If you experience flickering, set `use_popup=True`.\n",
						"\n",
						">**NOTE**: To use this notebook with a webcam, you need to run the notebook on a computer with a webcam. If you run the notebook on a remote server (for example, in Binder or Google Colab service), the webcam will not work. By default, the lower cell will run model inference on a video file. If you want to try live inference on your webcam set `WEBCAM_INFERENCE = True`"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "90708017",
					"metadata": {},
					"outputs": [],
					"source": [
						"WEBCAM_INFERENCE = False\n",
						"\n",
						"if WEBCAM_INFERENCE:\n",
						"    VIDEO_SOURCE = 0  # Webcam\n",
						"else:\n",
						"    if not Path(\"people.mp4\").exists():\n",
						"        download_file(\n",
						"            \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
						"            \"people.mp4\",\n",
						"        )\n",
						"    VIDEO_SOURCE = \"people.mp4\""
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "e8f1239c",
					"metadata": {},
					"outputs": [],
					"source": [
						"device"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"id": "440f2d0e",
					"metadata": {
						"tags": []
					},
					"outputs": [],
					"source": [
						"run_instance_segmentation(\n",
						"    source=VIDEO_SOURCE,\n",
						"    flip=True,\n",
						"    use_popup=False,\n",
						"    model=seg_ov_model,\n",
						"    device=device.value,\n",
						"    video_width=1280,\n",
						")"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
			{ modified: 24, original: 24 },
			{ modified: 25, original: 25 },
			{ modified: 26, original: 26 },
			{ modified: 27, original: 27 },
			{ modified: 28, original: 28 },
			{ modified: 29, original: 29 },
			{ modified: 30, original: 30 },
			{ modified: 31, original: 31 },
			{ modified: 32, original: 32 },
			{ modified: 33, original: 33 },
			{ modified: 34, original: 34 },
			{ modified: 35, original: 35 },
			{ modified: 36, original: 36 },
			{ modified: 37, original: 37 },
			{ modified: 38, original: 38 },
			{ modified: 39, original: 39 },
			{ modified: 40, original: 40 },
			{ modified: 41, original: 41 },
			{ modified: 42, original: 42 },
			{ modified: 43, original: 43 },
			{ modified: 44, original: 44 },
			{ modified: 45, original: 45 },
			{ modified: 46, original: 46 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 5, originalLength: 1, modifiedStart: 5, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 6, originalLength: 1, modifiedStart: 6, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 25, originalLength: 1, modifiedStart: 25, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 42, originalLength: 1, modifiedStart: 42, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 44, originalLength: 1, modifiedStart: 44, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 46, originalLength: 1, modifiedStart: 46, modifiedLength: 1 } satisfies IDiffChange,
		]);

	});
	test('Modification of a large cell and insert a new cell', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"id": "e9d010ff-85ba-4e69-b439-ba89e4a3a715",
					"source": [
						"# Convert and Optimize YOLOv10 with OpenVINO\n",
						"\n",
						"Real-time object detection aims to accurately predict object categories and positions in images with low latency. The YOLO series has been at the forefront of this research due to its balance between performance and efficiency. However, reliance on NMS and architectural inefficiencies have hindered optimal performance. YOLOv10 addresses these issues by introducing consistent dual assignments for NMS-free training and a holistic efficiency-accuracy driven model design strategy.\n",
						"\n",
						"YOLOv10, built on the [Ultralytics Python package](https://pypi.org/project/ultralytics/) by researchers at [Tsinghua University](https://www.tsinghua.edu.cn/en/), introduces a new approach to real-time object detection, addressing both the post-processing and model architecture deficiencies found in previous YOLO versions. By eliminating non-maximum suppression (NMS) and optimizing various model components, YOLOv10 achieves state-of-the-art performance with significantly reduced computational overhead. Extensive experiments demonstrate its superior accuracy-latency trade-offs across multiple model scales.\n",
						"\n",
						"![yolov10-approach.png](https://github.com/ultralytics/ultralytics/assets/26833433/f9b1bec0-928e-41ce-a205-e12db3c4929a)\n",
						"\n",
						"More details about model architecture you can find in original [repo](https://github.com/THU-MIG/yolov10), [paper](https://arxiv.org/abs/2405.14458) and [Ultralytics documentation](https://docs.ultralytics.com/models/yolov10/).\n",
						"\n",
						"This tutorial demonstrates step-by-step instructions on how to run and optimize PyTorch YOLO V10 with OpenVINO.\n",
						"\n",
						"The tutorial consists of the following steps:\n",
						"\n",
						"- Prepare PyTorch model\n",
						"- Convert PyTorch model to OpenVINO IR\n",
						"- Run model inference with OpenVINO\n",
						"- Prepare and run optimization pipeline using NNCF\n",
						"- Compare performance of the FP16 and quantized models.\n",
						"- Run optimized model inference on video\n",
						"- Launch interactive Gradio demo\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Download PyTorch model](#Download-PyTorch-model)\n",
						"- [Export PyTorch model to OpenVINO IR Format](#Export-PyTorch-model-to-OpenVINO-IR-Format)\n",
						"- [Run OpenVINO Inference on AUTO device using Ultralytics API](#Run-OpenVINO-Inference-on-AUTO-device-using-Ultralytics-API)\n",
						"- [Run OpenVINO Inference on selected device using Ultralytics API](#Run-OpenVINO-Inference-on-selected-device-using-Ultralytics-API)\n",
						"- [Optimize model using NNCF Post-training Quantization API](#Optimize-model-using-NNCF-Post-training-Quantization-API)\n",
						"    - [Prepare Quantization Dataset](#Prepare-Quantization-Dataset)\n",
						"    - [Quantize and Save INT8 model](#Quantize-and-Save-INT8-model)\n",
						"- [Run Optimized Model Inference](#Run-Optimized-Model-Inference)\n",
						"    - [Run Optimized Model on AUTO device](#Run-Optimized-Model-on-AUTO-device)\n",
						"    - [Run Optimized Model Inference on selected device](#Run-Optimized-Model-Inference-on-selected-device)\n",
						"- [Compare the Original and Quantized Models](#Compare-the-Original-and-Quantized-Models)\n",
						"    - [Model size](#Model-size)\n",
						"    - [Performance](#Performance)\n",
						"    - [FP16 model performance](#FP16-model-performance)\n",
						"    - [Int8 model performance](#Int8-model-performance)\n",
						"- [Live demo](#Live-demo)\n",
						"    - [Gradio Interactive Demo](#Gradio-Interactive-Demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/yolov10-optimization/yolov10-optimization.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"id": "144e86a1-c220-4e3b-a8b2-8eda443faf2d",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "bf76580b-860a-4dcd-8b01-f494cdffe37e",
					"source": [
						"import os\n",
						"\n",
						"os.environ[\"GIT_CLONE_PROTECTION_ACTIVE\"] = \"false\"\n",
						"\n",
						"%pip install -q \"nncf>=2.11.0\"\n",
						"%pip install -Uq \"openvino>=2024.3.0\"\n",
						"%pip install -q \"git+https://github.com/THU-MIG/yolov10.git\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
						"%pip install -q \"torch>=2.1\" \"torchvision>=0.16\" tqdm opencv-python \"gradio>=4.19\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "code",
					"id": "75772eac-565b-4234-82bf-f305e0c1ceda",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"from notebook_utils import download_file, VideoPlayer, device_widget"
					]
				},
				{
					"cell_type": "markdown",
					"id": "2888d8b4-44f6-4418-bae4-f433ff28010b",
					"source": [
						"## Download PyTorch model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"There are several version of [YOLO V10](https://github.com/THU-MIG/yolov10/tree/main?tab=readme-ov-file#performance) models provided by model authors. Each of them has different characteristics depends on number of training parameters, performance and accuracy. For demonstration purposes we will use `yolov10n`, but the same steps are also applicable to other models in YOLO V10 series."
					]
				},
				{
					"cell_type": "code",
					"id": "aea6315b-6a29-4f2a-96a9-eb3c1646a3a4",
					"source": [
						"models_dir = Path(\"./models\")\n",
						"models_dir.mkdir(exist_ok=True)"
					]
				},
				{
					"cell_type": "code",
					"id": "d828401e-b6bd-4796-a7b8-681957a159d4",
					"source": [
						"model_weights_url = \"https://github.com/jameslahm/yolov10/releases/download/v1.0/yolov10n.pt\"\n",
						"file_name = model_weights_url.split(\"/\")[-1]\n",
						"model_name = file_name.replace(\".pt\", \"\")\n",
						"\n",
						"download_file(model_weights_url, directory=models_dir)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "dc21487d-dd79-4b56-b8e2-5b28c8d92ac8",
					"source": [
						"## Export PyTorch model to OpenVINO IR Format\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "8cfec0d6-fb55-4611-861b-326e892fcf09",
					"source": [
						"As it was discussed before, YOLO V10 code is designed on top of [Ultralytics](https://docs.ultralytics.com/) library and has similar interface with YOLO V8 (You can check [YOLO V8 notebooks](https://github.com/openvinotoolkit/openvino_notebooks/tree/latest/notebooks/yolov8-optimization) for more detailed instruction how to work with Ultralytics API). Ultralytics support OpenVINO model export using [export](https://docs.ultralytics.com/modes/export/) method of model class. Additionally, we can specify parameters responsible for target input size, static or dynamic input shapes and model precision (FP32/FP16/INT8). INT8 quantization can be additionally performed on export stage, but for making approach more flexible, we consider how to perform quantization using [NNCF](https://github.com/openvinotoolkit/nncf)."
					]
				},
				{
					"cell_type": "code",
					"id": "d488f564-c1d5-4c6a-9b4a-353a3ab749e6",
					"source": [
						"import types\n",
						"from ultralytics.utils import ops, yaml_load, yaml_save\n",
						"from ultralytics import YOLOv10\n",
						"import torch\n",
						"\n",
						"detection_labels = {\n",
						"    0: \"person\",\n",
						"    1: \"bicycle\",\n",
						"    2: \"car\",\n",
						"    3: \"motorcycle\",\n",
						"    4: \"airplane\",\n",
						"    5: \"bus\",\n",
						"    6: \"train\",\n",
						"    7: \"truck\",\n",
						"    8: \"boat\",\n",
						"    9: \"traffic light\",\n",
						"    10: \"fire hydrant\",\n",
						"    11: \"stop sign\",\n",
						"    12: \"parking meter\",\n",
						"    13: \"bench\",\n",
						"    14: \"bird\",\n",
						"    15: \"cat\",\n",
						"    16: \"dog\",\n",
						"    17: \"horse\",\n",
						"    18: \"sheep\",\n",
						"    19: \"cow\",\n",
						"    20: \"elephant\",\n",
						"    21: \"bear\",\n",
						"    22: \"zebra\",\n",
						"    23: \"giraffe\",\n",
						"    24: \"backpack\",\n",
						"    25: \"umbrella\",\n",
						"    26: \"handbag\",\n",
						"    27: \"tie\",\n",
						"    28: \"suitcase\",\n",
						"    29: \"frisbee\",\n",
						"    30: \"skis\",\n",
						"    31: \"snowboard\",\n",
						"    32: \"sports ball\",\n",
						"    33: \"kite\",\n",
						"    34: \"baseball bat\",\n",
						"    35: \"baseball glove\",\n",
						"    36: \"skateboard\",\n",
						"    37: \"surfboard\",\n",
						"    38: \"tennis racket\",\n",
						"    39: \"bottle\",\n",
						"    40: \"wine glass\",\n",
						"    41: \"cup\",\n",
						"    42: \"fork\",\n",
						"    43: \"knife\",\n",
						"    44: \"spoon\",\n",
						"    45: \"bowl\",\n",
						"    46: \"banana\",\n",
						"    47: \"apple\",\n",
						"    48: \"sandwich\",\n",
						"    49: \"orange\",\n",
						"    50: \"broccoli\",\n",
						"    51: \"carrot\",\n",
						"    52: \"hot dog\",\n",
						"    53: \"pizza\",\n",
						"    54: \"donut\",\n",
						"    55: \"cake\",\n",
						"    56: \"chair\",\n",
						"    57: \"couch\",\n",
						"    58: \"potted plant\",\n",
						"    59: \"bed\",\n",
						"    60: \"dining table\",\n",
						"    61: \"toilet\",\n",
						"    62: \"tv\",\n",
						"    63: \"laptop\",\n",
						"    64: \"mouse\",\n",
						"    65: \"remote\",\n",
						"    66: \"keyboard\",\n",
						"    67: \"cell phone\",\n",
						"    68: \"microwave\",\n",
						"    69: \"oven\",\n",
						"    70: \"toaster\",\n",
						"    71: \"sink\",\n",
						"    72: \"refrigerator\",\n",
						"    73: \"book\",\n",
						"    74: \"clock\",\n",
						"    75: \"vase\",\n",
						"    76: \"scissors\",\n",
						"    77: \"teddy bear\",\n",
						"    78: \"hair drier\",\n",
						"    79: \"toothbrush\",\n",
						"}\n",
						"\n",
						"\n",
						"def v10_det_head_forward(self, x):\n",
						"    one2one = self.forward_feat([xi.detach() for xi in x], self.one2one_cv2, self.one2one_cv3)\n",
						"    if not self.export:\n",
						"        one2many = super().forward(x)\n",
						"\n",
						"    if not self.training:\n",
						"        one2one = self.inference(one2one)\n",
						"        if not self.export:\n",
						"            return {\"one2many\": one2many, \"one2one\": one2one}\n",
						"        else:\n",
						"            assert self.max_det != -1\n",
						"            boxes, scores, labels = ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det, self.nc)\n",
						"            return torch.cat(\n",
						"                [boxes, scores.unsqueeze(-1), labels.unsqueeze(-1).to(boxes.dtype)],\n",
						"                dim=-1,\n",
						"            )\n",
						"    else:\n",
						"        return {\"one2many\": one2many, \"one2one\": one2one}\n",
						"\n",
						"\n",
						"ov_model_path = models_dir / f\"{model_name}_openvino_model/{model_name}.xml\"\n",
						"if not ov_model_path.exists():\n",
						"    model = YOLOv10(models_dir / file_name)\n",
						"    model.model.model[-1].forward = types.MethodType(v10_det_head_forward, model.model.model[-1])\n",
						"    model.export(format=\"openvino\", dynamic=True, half=True)\n",
						"    config = yaml_load(ov_model_path.parent / \"metadata.yaml\")\n",
						"    config[\"names\"] = detection_labels\n",
						"    yaml_save(ov_model_path.parent / \"metadata.yaml\", config)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "8a23f621-f6cd-4bfd-8b98-24b8c0392761",
					"source": [
						"## Run OpenVINO Inference on AUTO device using Ultralytics API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"Now, when we exported model to OpenVINO, we can load it directly into YOLOv10 class, where automatic inference backend will provide easy-to-use user experience to run OpenVINO YOLOv10 model on the similar level like for original PyTorch model. The code bellow demonstrates how to run inference OpenVINO exported model with Ultralytics API on single image. [AUTO device](https://github.com/openvinotoolkit/openvino_notebooks/tree/latest/notebooks/auto-device) will be used for launching model."
					]
				},
				{
					"cell_type": "code",
					"id": "7e8ed361-bf09-4398-ba15-e6c5dda73cd3",
					"source": [
						"ov_yolo_model = YOLOv10(ov_model_path.parent, task=\"detect\")"
					]
				},
				{
					"cell_type": "code",
					"id": "8a473286-59b8-498f-8541-df84f041e119",
					"source": [
						"from PIL import Image\n",
						"\n",
						"IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
						"download_file(\n",
						"    url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
						"    filename=IMAGE_PATH.name,\n",
						"    directory=IMAGE_PATH.parent,\n",
						")"
					]
				},
				{
					"cell_type": "code",
					"id": "7e116e81-0ad4-4b9e-9f5c-58680fd8f0c6",
					"source": [
						"res = ov_yolo_model(IMAGE_PATH, iou=0.45, conf=0.2)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "943b9b04-4b97-4395-99d1-695465de1140",
					"source": [
						"## Run OpenVINO Inference on selected device using Ultralytics API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"In this part of notebook you can select inference device for running model inference to compare results with AUTO."
					]
				},
				{
					"cell_type": "code",
					"id": "41655bec-b006-4376-abbb-d6c8c09e89fb",
					"source": [
						"device = device_widget(\"CPU\")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "f284d5fd-12ba-4c08-afce-d8131ef7ad4d",
					"source": [
						"import openvino as ov\n",
						"\n",
						"core = ov.Core()\n",
						"\n",
						"ov_model = core.read_model(ov_model_path)\n",
						"\n",
						"# load model on selected device\n",
						"if \"GPU\" in device.value or \"NPU\" in device.value:\n",
						"    ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value:\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"det_compiled_model = core.compile_model(ov_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"id": "12164d21-948a-4ef7-9e1c-b0d41c53cf91",
					"source": [
						"ov_yolo_model.predictor.model.ov_compiled_model = det_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"id": "ecb9d396-ab1e-4398-a110-18c08b00ab14",
					"source": [
						"res = ov_yolo_model(IMAGE_PATH, iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "a4abf758-3a72-4ee0-afde-38076602f1c7",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "992b6ac3-4248-4848-8fd6-f9c1b9fc2051",
					"source": [
						"## Optimize model using NNCF Post-training Quantization API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
						"We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize YOLOv10.\n",
						"\n",
						"The optimization process contains the following steps:\n",
						"\n",
						"1. Create a Dataset for quantization.\n",
						"2. Run `nncf.quantize` for getting an optimized model.\n",
						"3. Serialize OpenVINO IR model, using the `openvino.save_model` function.\n",
						"\n",
						"Quantization is time and memory consuming process, you can skip this step using checkbox bellow:"
					]
				},
				{
					"cell_type": "code",
					"id": "ed9d9ce5-8df9-4803-ae85-bf0744de914c",
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"int8_model_det_path = models_dir / \"int8\" / f\"{model_name}_openvino_model/{model_name}.xml\"\n",
						"ov_yolo_int8_model = None\n",
						"\n",
						"to_quantize = widgets.Checkbox(\n",
						"    value=True,\n",
						"    description=\"Quantization\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"to_quantize"
					]
				},
				{
					"cell_type": "code",
					"id": "52287a23-b5ef-4c34-873a-32b398d26d1a",
					"source": [
						"# Fetch skip_kernel_extension module\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
						")\n",
						"open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
						"\n",
						"%load_ext skip_kernel_extension"
					]
				},
				{
					"cell_type": "markdown",
					"id": "89ba64af-aa73-4ed9-a2f9-527bd9ae8221",
					"source": [
						"### Prepare Quantization Dataset\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"For starting quantization, we need to prepare dataset. We will use validation subset from [MS COCO dataset](https://cocodataset.org/) for model quantization and Ultralytics validation data loader for preparing input data."
					]
				},
				{
					"cell_type": "code",
					"id": "5e10435e-9cd6-4f29-a865-721d6209314d",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"from zipfile import ZipFile\n",
						"\n",
						"from ultralytics.data.utils import DATASETS_DIR\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"\n",
						"    DATA_URL = \"http://images.cocodataset.org/zips/val2017.zip\"\n",
						"    LABELS_URL = \"https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels-segments.zip\"\n",
						"    CFG_URL = \"https://raw.githubusercontent.com/ultralytics/ultralytics/v8.1.0/ultralytics/cfg/datasets/coco.yaml\"\n",
						"\n",
						"    OUT_DIR = DATASETS_DIR\n",
						"\n",
						"    DATA_PATH = OUT_DIR / \"val2017.zip\"\n",
						"    LABELS_PATH = OUT_DIR / \"coco2017labels-segments.zip\"\n",
						"    CFG_PATH = OUT_DIR / \"coco.yaml\"\n",
						"\n",
						"    download_file(DATA_URL, DATA_PATH.name, DATA_PATH.parent)\n",
						"    download_file(LABELS_URL, LABELS_PATH.name, LABELS_PATH.parent)\n",
						"    download_file(CFG_URL, CFG_PATH.name, CFG_PATH.parent)\n",
						"\n",
						"    if not (OUT_DIR / \"coco/labels\").exists():\n",
						"        with ZipFile(LABELS_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR)\n",
						"        with ZipFile(DATA_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR / \"coco/images\")"
					]
				},
				{
					"cell_type": "code",
					"id": "b0a14818-21fd-4c71-8e58-2297dfafaadd",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"from ultralytics.utils import DEFAULT_CFG\n",
						"from ultralytics.cfg import get_cfg\n",
						"from ultralytics.data.converter import coco80_to_coco91_class\n",
						"from ultralytics.data.utils import check_det_dataset\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    args = get_cfg(cfg=DEFAULT_CFG)\n",
						"    args.data = str(CFG_PATH)\n",
						"    det_validator = ov_yolo_model.task_map[ov_yolo_model.task][\"validator\"](args=args)\n",
						"\n",
						"    det_validator.data = check_det_dataset(args.data)\n",
						"    det_validator.stride = 32\n",
						"    det_data_loader = det_validator.get_dataloader(OUT_DIR / \"coco\", 1)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "83fa49f4-e40c-48c7-82a6-8ef56e20b343",
					"source": [
						"NNCF provides `nncf.Dataset` wrapper for using native framework dataloaders in quantization pipeline. Additionally, we specify transform function that will be responsible for preparing input data in model expected format."
					]
				},
				{
					"cell_type": "code",
					"id": "253f0e0f-131e-43b9-b1be-639d4b3d1fe5",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"import nncf\n",
						"from typing import Dict\n",
						"\n",
						"\n",
						"def transform_fn(data_item:Dict):\n",
						"    \"\"\"\n",
						"    Quantization transform function. Extracts and preprocess input data from dataloader item for quantization.\n",
						"    Parameters:\n",
						"       data_item: Dict with data item produced by DataLoader during iteration\n",
						"    Returns:\n",
						"        input_tensor: Input data for quantization\n",
						"    \"\"\"\n",
						"    input_tensor = det_validator.preprocess(data_item)['img'].numpy()\n",
						"    return input_tensor\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    quantization_dataset = nncf.Dataset(det_data_loader, transform_fn)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "ff8a6372-2097-4308-a4e5-e4651242a8df",
					"source": [
						"### Quantize and Save INT8 model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The `nncf.quantize` function provides an interface for model quantization. It requires an instance of the OpenVINO Model and quantization dataset. \n",
						"Optionally, some additional parameters for the configuration quantization process (number of samples for quantization, preset, ignored scope, etc.) can be provided. YOLOv10 model contains non-ReLU activation functions, which require asymmetric quantization of activations. To achieve a better result, we will use a `mixed` quantization preset. It provides symmetric quantization of weights and asymmetric quantization of activations.\n",
						"\n",
						">**Note**: Model post-training quantization is time-consuming process. Be patient, it can take several minutes depending on your hardware."
					]
				},
				{
					"cell_type": "code",
					"id": "4c12ae38-239e-4fe1-8296-47567f98538c",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"import shutil\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    quantized_det_model = nncf.quantize(\n",
						"        ov_model,\n",
						"        quantization_dataset,\n",
						"        preset=nncf.QuantizationPreset.MIXED,\n",
						"    )\n",
						"\n",
						"    ov.save_model(quantized_det_model,  int8_model_det_path)\n",
						"    shutil.copy(ov_model_path.parent / \"metadata.yaml\", int8_model_det_path.parent / \"metadata.yaml\")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "41b8a3eb-adb8-40cf-8602-311ff8d75a76",
					"source": [
						"## Run Optimized Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The way of usage INT8 quantized model is the same like for model before quantization. Let's check inference result of quantized model on single image"
					]
				},
				{
					"cell_type": "markdown",
					"id": "f438df41-d614-4cca-a746-a173d92dddf5",
					"source": [
						"### Run Optimized Model on AUTO device\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "af61303b-16bd-4206-b926-2c28b3a859a4",
					"source": [
						"%%skip not $to_quantize.value\n",
						"ov_yolo_int8_model = YOLOv10(int8_model_det_path.parent, task=\"detect\")"
					]
				},
				{
					"cell_type": "code",
					"id": "db13b8dd-49a1-4046-8a1a-491eee095a1b",
					"source": [
						"%%skip not $to_quantize.value\n",
						"res = ov_yolo_int8_model(IMAGE_PATH, iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "0423ba31-f0ee-4bfc-a18a-f7b870a9683d",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "4b4f7118-e046-4f87-b403-87053f2b0ac3",
					"source": [
						"### Run Optimized Model Inference on selected device\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "142f22ae-6a8d-4578-b97f-4d9c77123418",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "2f43767a-1994-4202-a411-4738c92223da",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value or \"NPU\" in device.value:\n",
						"    ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value:\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"\n",
						"quantized_det_model = core.read_model(int8_model_det_path)\n",
						"quantized_det_compiled_model = core.compile_model(quantized_det_model, device.value, ov_config)\n",
						"\n",
						"ov_yolo_int8_model.predictor.model.ov_compiled_model = quantized_det_compiled_model\n",
						"\n",
						"res = ov_yolo_int8_model(IMAGE_PATH,  iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "51d64c7b-7595-41ca-87dd-4c85308f84d6",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "c7b70f75-a706-4f24-a487-88a3ad645dd5",
					"source": [
						"## Compare the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "d6561246-a9e0-4cdf-8207-7e2f919d8582",
					"source": [
						"### Model size\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "43d84d15-75be-4430-a58b-61cfd8e08dd0",
					"source": [
						"ov_model_weights = ov_model_path.with_suffix(\".bin\")\n",
						"print(f\"Size of FP16 model is {ov_model_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
						"if int8_model_det_path.exists():\n",
						"    ov_int8_weights = int8_model_det_path.with_suffix(\".bin\")\n",
						"    print(f\"Size of model with INT8 compressed weights is {ov_int8_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
						"    print(f\"Compression rate for INT8 model: {ov_model_weights.stat().st_size / ov_int8_weights.stat().st_size:.3f}\")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "f72b3d84-91f3-493c-8f4e-46598ffbd6d1",
					"source": [
						"### Performance\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"### FP16 model performance\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "909ea508-e007-4313-ab0f-bfa0e3906176",
					"source": [
						"!benchmark_app -m $ov_model_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "markdown",
					"id": "9852a8ef-4c45-43dd-a9ef-53ca6f13c99e",
					"source": [
						"### Int8 model performance\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "8254626d-6c4f-4ca3-936f-a7c47e402e03",
					"source": [
						"if int8_model_det_path.exists():\n",
						"    !benchmark_app -m $int8_model_det_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "markdown",
					"id": "12d35bc1-8101-45ea-8bd7-53720b98f5e5",
					"source": [
						"## Live demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The following code runs model inference on a video:"
					]
				},
				{
					"cell_type": "code",
					"id": "246bd379-b1d5-4eb2-928f-97c7c2455a92",
					"source": [
						"import collections\n",
						"import time\n",
						"from IPython import display\n",
						"import cv2\n",
						"import numpy as np\n",
						"\n",
						"\n",
						"# Main processing function to run object detection.\n",
						"def run_object_detection(\n",
						"    source=0,\n",
						"    flip=False,\n",
						"    use_popup=False,\n",
						"    skip_first_frames=0,\n",
						"    det_model=ov_yolo_int8_model,\n",
						"    device=device.value,\n",
						"):\n",
						"    player = None\n",
						"    try:\n",
						"        # Create a video player to play with target fps.\n",
						"        player = VideoPlayer(source=source, flip=flip, fps=30, skip_first_frames=skip_first_frames)\n",
						"        # Start capturing.\n",
						"        player.start()\n",
						"        if use_popup:\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(winname=title, flags=cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"        while True:\n",
						"            # Grab the frame.\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"            # If the frame is larger than full HD, reduce size to improve the performance.\n",
						"            scale = 1280 / max(frame.shape)\n",
						"            if scale < 1:\n",
						"                frame = cv2.resize(\n",
						"                    src=frame,\n",
						"                    dsize=None,\n",
						"                    fx=scale,\n",
						"                    fy=scale,\n",
						"                    interpolation=cv2.INTER_AREA,\n",
						"                )\n",
						"            # Get the results.\n",
						"            input_image = np.array(frame)\n",
						"\n",
						"            start_time = time.time()\n",
						"            detections = det_model(input_image, iou=0.45, conf=0.2, verbose=False)\n",
						"            stop_time = time.time()\n",
						"            frame = detections[0].plot()\n",
						"\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Use processing times from last 200 frames.\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            _, f_width = frame.shape[:2]\n",
						"            # Mean processing time [ms].\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"            cv2.putText(\n",
						"                img=frame,\n",
						"                text=f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                org=(20, 40),\n",
						"                fontFace=cv2.FONT_HERSHEY_COMPLEX,\n",
						"                fontScale=f_width / 1000,\n",
						"                color=(0, 0, 255),\n",
						"                thickness=1,\n",
						"                lineType=cv2.LINE_AA,\n",
						"            )\n",
						"            # Use this workaround if there is flickering.\n",
						"            if use_popup:\n",
						"                cv2.imshow(winname=title, mat=frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # Encode numpy array to jpg.\n",
						"                _, encoded_img = cv2.imencode(ext=\".jpg\", img=frame, params=[cv2.IMWRITE_JPEG_QUALITY, 100])\n",
						"                # Create an IPython image.\n",
						"                i = display.Image(data=encoded_img)\n",
						"                # Display the image in this notebook.\n",
						"                display.clear_output(wait=True)\n",
						"                display.display(i)\n",
						"    # ctrl-c\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    # any different error\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        if player is not None:\n",
						"            # Stop capturing.\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()"
					]
				},
				{
					"cell_type": "code",
					"id": "8db930f4-fc41-4635-9248-deab2b6cfcb8",
					"source": [
						"use_int8 = widgets.Checkbox(\n",
						"    value=ov_yolo_int8_model is not None,\n",
						"    description=\"Use int8 model\",\n",
						"    disabled=ov_yolo_int8_model is None,\n",
						")\n",
						"\n",
						"use_int8"
					]
				},
				{
					"cell_type": "code",
					"id": "59e37135-1bdc-40ff-8789-b5b4491cab84",
					"source": [
						"WEBCAM_INFERENCE = False\n",
						"\n",
						"if WEBCAM_INFERENCE:\n",
						"    VIDEO_SOURCE = 0  # Webcam\n",
						"else:\n",
						"    download_file(\n",
						"        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
						"        directory=\"data\",\n",
						"    )\n",
						"    VIDEO_SOURCE = \"data/people.mp4\""
					]
				},
				{
					"cell_type": "code",
					"id": "88c6f47c-1c47-451d-89ba-0a29760f5ec4",
					"source": [
						"run_object_detection(\n",
						"    det_model=ov_yolo_model if not use_int8.value else ov_yolo_int8_model,\n",
						"    source=VIDEO_SOURCE,\n",
						"    flip=True,\n",
						"    use_popup=False,\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "9f612d80-68e2-45f8-aa62-9f3a05ca9de8",
					"source": [
						"### Gradio Interactive Demo\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "ffdbe248-7086-4804-b4b9-c3dd454bf80c",
					"source": [
						"import gradio as gr\n",
						"\n",
						"\n",
						"def yolov10_inference(image, int8, conf_threshold, iou_threshold):\n",
						"    model = ov_yolo_model if not int8 else ov_yolo_int8_model\n",
						"    results = model(source=image, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]\n",
						"    annotated_image = Image.fromarray(results.plot())\n",
						"\n",
						"    return annotated_image\n",
						"\n",
						"\n",
						"with gr.Blocks() as demo:\n",
						"    gr.HTML(\n",
						"        \"\"\"\n",
						"    <h1 style='text-align: center'>\n",
						"    YOLOv10: Real-Time End-to-End Object Detection using OpenVINO\n",
						"    </h1>\n",
						"    \"\"\"\n",
						"    )\n",
						"    with gr.Row():\n",
						"        with gr.Column():\n",
						"            image = gr.Image(type=\"numpy\", label=\"Image\")\n",
						"            conf_threshold = gr.Slider(\n",
						"                label=\"Confidence Threshold\",\n",
						"                minimum=0.1,\n",
						"                maximum=1.0,\n",
						"                step=0.1,\n",
						"                value=0.2,\n",
						"            )\n",
						"            iou_threshold = gr.Slider(\n",
						"                label=\"IoU Threshold\",\n",
						"                minimum=0.1,\n",
						"                maximum=1.0,\n",
						"                step=0.1,\n",
						"                value=0.45,\n",
						"            )\n",
						"            use_int8 = gr.Checkbox(\n",
						"                value=ov_yolo_int8_model is not None,\n",
						"                visible=ov_yolo_int8_model is not None,\n",
						"                label=\"Use INT8 model\",\n",
						"            )\n",
						"            yolov10_infer = gr.Button(value=\"Detect Objects\")\n",
						"\n",
						"        with gr.Column():\n",
						"            output_image = gr.Image(type=\"pil\", label=\"Annotated Image\")\n",
						"\n",
						"        yolov10_infer.click(\n",
						"            fn=yolov10_inference,\n",
						"            inputs=[\n",
						"                image,\n",
						"                use_int8,\n",
						"                conf_threshold,\n",
						"                iou_threshold,\n",
						"            ],\n",
						"            outputs=[output_image],\n",
						"        )\n",
						"    examples = gr.Examples(\n",
						"        [\n",
						"            \"data/coco_bike.jpg\",\n",
						"        ],\n",
						"        inputs=[\n",
						"            image,\n",
						"        ],\n",
						"    )\n",
						"\n",
						"\n",
						"try:\n",
						"    demo.launch(debug=True)\n",
						"except Exception:\n",
						"    demo.launch(debug=True, share=True)"
					]
				}
			]
				.map(fromJupyterCell)
			,
			[
				{
					"cell_type": "markdown",
					"id": "e9d010ff-85ba-4e69-b439-ba89e4a3a715",
					"source": [
						"# Convert and Optimize YOLOv10 with OpenVINO\n",
						"\n",
						"Real-time object detection aims to accurately predict object categories and positions in images with low latency. The YOLO series has been at the forefront of this research due to its balance between performance and efficiency. However, reliance on NMS and architectural inefficiencies have hindered optimal performance. YOLOv10 addresses these issues by introducing consistent dual assignments for NMS-free training and a holistic efficiency-accuracy driven model design strategy.\n",
						"\n",
						"YOLOv10, built on the [Ultralytics Python package](https://pypi.org/project/ultralytics/) by researchers at [Tsinghua University](https://www.tsinghua.edu.cn/en/), introduces a new approach to real-time object detection, addressing both the post-processing and model architecture deficiencies found in previous YOLO versions. By eliminating non-maximum suppression (NMS) and optimizing various model components, YOLOv10 achieves state-of-the-art performance with significantly reduced computational overhead. Extensive experiments demonstrate its superior accuracy-latency trade-offs across multiple model scales.\n",
						"\n",
						"![yolov10-approach.png](https://github.com/ultralytics/ultralytics/assets/26833433/f9b1bec0-928e-41ce-a205-e12db3c4929a)\n",
						"\n",
						"More details about model architecture you can find in original [repo](https://github.com/THU-MIG/yolov10), [paper](https://arxiv.org/abs/2405.14458) and [Ultralytics documentation](https://docs.ultralytics.com/models/yolov10/).\n",
						"\n",
						"This tutorial demonstrates step-by-step instructions on how to run and optimize PyTorch YOLO V10 with OpenVINO.\n",
						"\n",
						"The tutorial consists of the following steps:\n",
						"\n",
						"- Prepare PyTorch model\n",
						"- Convert PyTorch model to OpenVINO IR\n",
						"- Run model inference with OpenVINO\n",
						"- Prepare and run optimization pipeline using NNCF\n",
						"- Compare performance of the FP16 and quantized models.\n",
						"- Run optimized model inference on video\n",
						"- Launch interactive Gradio demo\n",
						"\n",
						"#### Table of contents:\n",
						"\n",
						"- [Prerequisites](#Prerequisites)\n",
						"- [Download PyTorch model](#Download-PyTorch-model)\n",
						"- [Export PyTorch model to OpenVINO IR Format](#Export-PyTorch-model-to-OpenVINO-IR-Format)\n",
						"- [Run OpenVINO Inference on AUTO device using Ultralytics API](#Run-OpenVINO-Inference-on-AUTO-device-using-Ultralytics-API)\n",
						"- [Run OpenVINO Inference on selected device using Ultralytics API](#Run-OpenVINO-Inference-on-selected-device-using-Ultralytics-API)\n",
						"- [Optimize model using NNCF Post-training Quantization API](#Optimize-model-using-NNCF-Post-training-Quantization-API)\n",
						"    - [Prepare Quantization Dataset](#Prepare-Quantization-Dataset)\n",
						"    - [Quantize and Save INT8 model](#Quantize-and-Save-INT8-model)\n",
						"- [Run Optimized Model Inference](#Run-Optimized-Model-Inference)\n",
						"    - [Run Optimized Model on AUTO device](#Run-Optimized-Model-on-AUTO-device)\n",
						"    - [Run Optimized Model Inference on selected device](#Run-Optimized-Model-Inference-on-selected-device)\n",
						"- [Compare the Original and Quantized Models](#Compare-the-Original-and-Quantized-Models)\n",
						"    - [Model size](#Model-size)\n",
						"    - [Performance](#Performance)\n",
						"    - [FP16 model performance](#FP16-model-performance)\n",
						"    - [Int8 model performance](#Int8-model-performance)\n",
						"- [Live demo](#Live-demo)\n",
						"    - [Gradio Interactive Demo](#Gradio-Interactive-Demo)\n",
						"\n",
						"\n",
						"### Installation Instructions\n",
						"\n",
						"This is a self-contained example that relies solely on its own code.\n",
						"\n",
						"We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
						"For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
						"\n",
						"<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/yolov10-optimization/yolov10-optimization.ipynb\" />\n"
					]
				},
				{
					"cell_type": "markdown",
					"id": "144e86a1-c220-4e3b-a8b2-8eda443faf2d",
					"source": [
						"## Prerequisites\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "bf76580b-860a-4dcd-8b01-f494cdffe37e",
					"source": [
						"import os\n",
						"\n",
						"os.environ[\"GIT_CLONE_PROTECTION_ACTIVE\"] = \"false\"\n",
						"\n",
						"%pip install -q \"nncf>=2.11.0\"\n",
						"%pip install -Uq \"openvino>=2024.3.0\"\n",
						"%pip install -q \"git+https://github.com/THU-MIG/yolov10.git\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
						"%pip install -q \"torch>=2.1\" \"torchvision>=0.16\" tqdm opencv-python \"gradio>=4.19\" --extra-index-url https://download.pytorch.org/whl/cpu"
					]
				},
				{
					"cell_type": "code",
					"id": "75772eac-565b-4234-82bf-f305e0c1ceda",
					"source": [
						"from pathlib import Path\n",
						"\n",
						"# Fetch `notebook_utils` module\n",
						"import requests\n",
						"\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
						")\n",
						"\n",
						"open(\"notebook_utils.py\", \"w\").write(r.text)\n",
						"\n",
						"from notebook_utils import download_file, VideoPlayer, device_widget"
					]
				},
				{
					"cell_type": "markdown",
					"id": "2888d8b4-44f6-4418-bae4-f433ff28010b",
					"source": [
						"## Download PyTorch model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"There are several version of [YOLO V10](https://github.com/THU-MIG/yolov10/tree/main?tab=readme-ov-file#performance) models provided by model authors. Each of them has different characteristics depends on number of training parameters, performance and accuracy. For demonstration purposes we will use `yolov10n`, but the same steps are also applicable to other models in YOLO V10 series."
					]
				},
				{
					"cell_type": "code",
					"id": "aea6315b-6a29-4f2a-96a9-eb3c1646a3a4",
					"source": [
						"models_dir = Path(\"./models\")\n",
						"models_dir.mkdir(exist_ok=True)"
					]
				},
				{
					"cell_type": "code",
					"id": "d828401e-b6bd-4796-a7b8-681957a159d4",
					"source": [
						"model_weights_url = \"https://github.com/jameslahm/yolov10/releases/download/v1.0/yolov10n.pt\"\n",
						"file_name = model_weights_url.split(\"/\")[-1]\n",
						"model_name = file_name.replace(\".pt\", \"\")\n",
						"\n",
						"download_file(model_weights_url, directory=models_dir)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "dc21487d-dd79-4b56-b8e2-5b28c8d92ac8",
					"source": [
						"## Export PyTorch model to OpenVINO IR Format\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "8cfec0d6-fb55-4611-861b-326e892fcf09",
					"source": [
						"As it was discussed before, YOLO V10 code is designed on top of [Ultralytics](https://docs.ultralytics.com/) library and has similar interface with YOLO V8 (You can check [YOLO V8 notebooks](https://github.com/openvinotoolkit/openvino_notebooks/tree/latest/notebooks/yolov8-optimization) for more detailed instruction how to work with Ultralytics API). Ultralytics support OpenVINO model export using [export](https://docs.ultralytics.com/modes/export/) method of model class. Additionally, we can specify parameters responsible for target input size, static or dynamic input shapes and model precision (FP32/FP16/INT8). INT8 quantization can be additionally performed on export stage, but for making approach more flexible, we consider how to perform quantization using [NNCF](https://github.com/openvinotoolkit/nncf)."
					]
				},
				{
					"cell_type": "code",
					"id": "d488f564-c1d5-4c6a-9b4a-353a3ab749e6",
					"source": [
						"import types\n",
						"from ultralytics.utils import ops, yaml_load, yaml_save\n",
						"from ultralytics import YOLOv10\n",
						"import torch\n",
						"\n",
						"detection_labels = {\n",
						"    0: \"person\",\n",
						"    1: \"bicycle\",\n",
						"    2: \"car\",\n",
						"    3: \"motorcycle\",\n",
						"    4: \"airplane\",\n",
						"    5: \"bus\",\n",
						"    6: \"train\",\n",
						"    7: \"truck\",\n",
						"    8: \"boat\",\n",
						"    9: \"traffic light\",\n",
						"    10: \"fire hydrant\",\n",
						"    11: \"stop sign\",\n",
						"    12: \"parking meter\",\n",
						"    13: \"bench\",\n",
						"    14: \"bird\",\n",
						"    15: \"cat\",\n",
						"    16: \"dog\",\n",
						"    17: \"horse\",\n",
						"    18: \"sheep\",\n",
						"    19: \"cow\",\n",
						"    20: \"elephant\",\n",
						"    21: \"bear\",\n",
						"    22: \"zebra\",\n",
						"    23: \"giraffe\",\n",
						"    24: \"backpack\",\n",
						"    25: \"umbrella\",\n",
						"    26: \"handbag\",\n",
						"    27: \"tie\",\n",
						"    28: \"suitcase\",\n",
						"    29: \"frisbee\",\n",
						"    30: \"skis\",\n",
						"    31: \"snowboard\",\n",
						"    32: \"sports ball\",\n",
						"    33: \"kite\",\n",
						"    34: \"baseball bat\",\n",
						"    35: \"baseball glove\",\n",
						"    36: \"skateboard\",\n",
						"    37: \"surfboard\",\n",
						"    38: \"tennis racket\",\n",
						"    39: \"bottle\",\n",
						"    40: \"wine glass\",\n",
						"    41: \"cup\",\n",
						"    42: \"fork\",\n",
						"    43: \"knife\",\n",
						"    44: \"spoon\",\n",
						"    45: \"bowl\",\n",
						"    46: \"banana\",\n",
						"    47: \"apple\",\n",
						"    48: \"sandwich\",\n",
						"    49: \"orange\",\n",
						"    50: \"broccoli\",\n",
						"    51: \"carrot\",\n",
						"    52: \"hot dog\",\n",
						"    53: \"pizza\",\n",
						"    54: \"donut\",\n",
						"    55: \"cake\",\n",
						"    56: \"chair\",\n",
						"    57: \"couch\",\n",
						"    58: \"potted plant\",\n",
						"    59: \"bed\",\n",
						"    60: \"dining table\",\n",
						"    61: \"toilet\",\n",
						"    62: \"tv\",\n",
						"    63: \"laptop\",\n",
						"    64: \"mouse\",\n",
						"    65: \"remote\",\n",
						"    66: \"keyboard\",\n",
						"    67: \"cell phone\",\n",
						"    68: \"microwave\",\n",
						"    69: \"oven\",\n",
						"    70: \"toaster\",\n",
						"    71: \"sink\",\n",
						"    72: \"refrigerator\",\n",
						"    73: \"book\",\n",
						"    74: \"clock\",\n",
						"    75: \"vase\",\n",
						"    76: \"scissors\",\n",
						"    77: \"teddy bear\",\n",
						"    78: \"hair drier\",\n",
						"    79: \"toothbrush\",\n",
						"}\n",
						"\n",
						"\n",
						"def v10_det_head_forward(self, x):\n",
						"    one2one = self.forward_feat([xi.detach() for xi in x], self.one2one_cv2, self.one2one_cv3)\n",
						"    if not self.export:\n",
						"        one2many = super().forward(x)\n",
						"\n",
						"    if not self.training:\n",
						"        one2one = self.inference(one2one)\n",
						"        if not self.export:\n",
						"            return {\"one2many\": one2many, \"one2one\": one2one}\n",
						"        else:\n",
						"            assert self.max_det != -1\n",
						"            boxes, scores, labels = ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det, self.nc)\n",
						"            return torch.cat(\n",
						"                [boxes, scores.unsqueeze(-1), labels.unsqueeze(-1).to(boxes.dtype)],\n",
						"                dim=-1,\n",
						"            )\n",
						"    else:\n",
						"        return {\"one2many\": one2many, \"one2one\": one2one}\n",
						"\n",
						"\n",
						"ov_model_path = models_dir / f\"{model_name}_openvino_model/{model_name}.xml\"\n",
						"if not ov_model_path.exists():\n",
						"    model = YOLOv10(models_dir / file_name)\n",
						"    model.model.model[-1].forward = types.MethodType(v10_det_head_forward, model.model.model[-1])\n",
						"    model.export(format=\"openvino\", dynamic=True, half=True)\n",
						"    config = yaml_load(ov_model_path.parent / \"metadata.yaml\")\n",
						"    config[\"names\"] = detection_labels\n",
						"    yaml_save(ov_model_path.parent / \"metadata.yaml\", config)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "8a23f621-f6cd-4bfd-8b98-24b8c0392761",
					"source": [
						"## Run OpenVINO Inference on AUTO device using Ultralytics API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"Now, when we exported model to OpenVINO, we can load it directly into YOLOv10 class, where automatic inference backend will provide easy-to-use user experience to run OpenVINO YOLOv10 model on the similar level like for original PyTorch model. The code bellow demonstrates how to run inference OpenVINO exported model with Ultralytics API on single image. [AUTO device](https://github.com/openvinotoolkit/openvino_notebooks/tree/latest/notebooks/auto-device) will be used for launching model."
					]
				},
				{
					"cell_type": "code",
					"id": "7e8ed361-bf09-4398-ba15-e6c5dda73cd3",
					"source": [
						"ov_yolo_model = YOLOv10(ov_model_path.parent, task=\"detect\")"
					]
				},
				{
					"cell_type": "code",
					"id": "8a473286-59b8-498f-8541-df84f041e119",
					"source": [
						"from PIL import Image\n",
						"\n",
						"IMAGE_PATH = Path(\"./data/coco_bike.jpg\")\n",
						"download_file(\n",
						"    url=\"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/coco_bike.jpg\",\n",
						"    filename=IMAGE_PATH.name,\n",
						"    directory=IMAGE_PATH.parent,\n",
						")"
					]
				},
				{
					"cell_type": "code",
					"id": "7e116e81-0ad4-4b9e-9f5c-58680fd8f0c6",
					"source": [
						"res = ov_yolo_model(IMAGE_PATH, iou=0.45, conf=0.2)\n",
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "943b9b04-4b97-4395-99d1-695465de1140",
					"source": [
						"## Run OpenVINO Inference on selected device using Ultralytics API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"In this part of notebook you can select inference device for running model inference to compare results with AUTO."
					]
				},
				{
					"cell_type": "code",
					"id": "41655bec-b006-4376-abbb-d6c8c09e89fb",
					"source": [
						"device = device_widget(\"CPU\")\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "f284d5fd-12ba-4c08-afce-d8131ef7ad4d",
					"source": [
						"import openvino as ov\n",
						"\n",
						"core = ov.Core()\n",
						"\n",
						"ov_model = core.read_model(ov_model_path)\n",
						"\n",
						"# load model on selected device\n",
						"if \"GPU\" in device.value or \"NPU\" in device.value:\n",
						"    ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value:\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"det_compiled_model = core.compile_model(ov_model, device.value, ov_config)"
					]
				},
				{
					"cell_type": "code",
					"id": "12164d21-948a-4ef7-9e1c-b0d41c53cf91",
					"source": [
						"ov_yolo_model.predictor.model.ov_compiled_model = det_compiled_model"
					]
				},
				{
					"cell_type": "code",
					"id": "ecb9d396-ab1e-4398-a110-18c08b00ab14",
					"source": [
						"res = ov_yolo_model(IMAGE_PATH, iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "a4abf758-3a72-4ee0-afde-38076602f1c7",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "992b6ac3-4248-4848-8fd6-f9c1b9fc2051",
					"source": [
						"## Optimize model using NNCF Post-training Quantization API\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
						"We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize YOLOv10.\n",
						"\n",
						"The optimization process contains the following steps:\n",
						"\n",
						"1. Create a Dataset for quantization.\n",
						"2. Run `nncf.quantize` for getting an optimized model.\n",
						"3. Serialize OpenVINO IR model, using the `openvino.save_model` function.\n",
						"\n",
						"Quantization is time and memory consuming process, you can skip this step using checkbox bellow:"
					]
				},
				{
					"cell_type": "code",
					"id": "ed9d9ce5-8df9-4803-ae85-bf0744de914c",
					"source": [
						"import ipywidgets as widgets\n",
						"\n",
						"int8_model_det_path = models_dir / \"int8\" / f\"{model_name}_openvino_model/{model_name}.xml\"\n",
						"ov_yolo_int8_model = None\n",
						"\n",
						"to_quantize = widgets.Checkbox(\n",
						"    value=True,\n",
						"    description=\"Quantization\",\n",
						"    disabled=False,\n",
						")\n",
						"\n",
						"to_quantize"
					]
				},
				{
					"cell_type": "code",
					"id": "52287a23-b5ef-4c34-873a-32b398d26d1a",
					"source": [
						"# Fetch skip_kernel_extension module\n",
						"r = requests.get(\n",
						"    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
						")\n",
						"open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
						"\n",
						"%load_ext skip_kernel_extension"
					]
				},
				{
					"cell_type": "markdown",
					"id": "89ba64af-aa73-4ed9-a2f9-527bd9ae8221",
					"source": [
						"### Prepare Quantization Dataset\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"For starting quantization, we need to prepare dataset. We will use validation subset from [MS COCO dataset](https://cocodataset.org/) for model quantization and Ultralytics validation data loader for preparing input data."
					]
				},
				{
					"cell_type": "code",
					"id": "5e10435e-9cd6-4f29-a865-721d6209314d",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"from zipfile import ZipFile\n",
						"\n",
						"from ultralytics.data.utils import DATASETS_DIR\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"\n",
						"    DATA_URL = \"http://images.cocodataset.org/zips/val2017.zip\"\n",
						"    LABELS_URL = \"https://github.com/ultralytics/yolov5/releases/download/v1.0/coco2017labels-segments.zip\"\n",
						"    CFG_URL = \"https://raw.githubusercontent.com/ultralytics/ultralytics/v8.1.0/ultralytics/cfg/datasets/coco.yaml\"\n",
						"\n",
						"    OUT_DIR = DATASETS_DIR\n",
						"\n",
						"    DATA_PATH = OUT_DIR / \"val2017.zip\"\n",
						"    LABELS_PATH = OUT_DIR / \"coco2017labels-segments.zip\"\n",
						"    CFG_PATH = OUT_DIR / \"coco.yaml\"\n",
						"\n",
						"    download_file(DATA_URL, DATA_PATH.name, DATA_PATH.parent)\n",
						"    download_file(LABELS_URL, LABELS_PATH.name, LABELS_PATH.parent)\n",
						"    download_file(CFG_URL, CFG_PATH.name, CFG_PATH.parent)\n",
						"\n",
						"    if not (OUT_DIR / \"coco/labels\").exists():\n",
						"        with ZipFile(LABELS_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR)\n",
						"        with ZipFile(DATA_PATH, \"r\") as zip_ref:\n",
						"            zip_ref.extractall(OUT_DIR / \"coco/images\")"
					]
				},
				{
					"cell_type": "code",
					"id": "b0a14818-21fd-4c71-8e58-2297dfafaadd",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"from ultralytics.utils import DEFAULT_CFG\n",
						"from ultralytics.cfg import get_cfg\n",
						"from ultralytics.data.converter import coco80_to_coco91_class\n",
						"from ultralytics.data.utils import check_det_dataset\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    args = get_cfg(cfg=DEFAULT_CFG)\n",
						"    args.data = str(CFG_PATH)\n",
						"    det_validator = ov_yolo_model.task_map[ov_yolo_model.task][\"validator\"](args=args)\n",
						"\n",
						"    det_validator.data = check_det_dataset(args.data)\n",
						"    det_validator.stride = 32\n",
						"    det_data_loader = det_validator.get_dataloader(OUT_DIR / \"coco\", 1)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "83fa49f4-e40c-48c7-82a6-8ef56e20b343",
					"source": [
						"NNCF provides `nncf.Dataset` wrapper for using native framework dataloaders in quantization pipeline. Additionally, we specify transform function that will be responsible for preparing input data in model expected format."
					]
				},
				{
					"cell_type": "code",
					"id": "253f0e0f-131e-43b9-b1be-639d4b3d1fe5",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"import nncf\n",
						"from typing import Dict\n",
						"\n",
						"\n",
						"def transform_fn(data_item:Dict):\n",
						"    \"\"\"\n",
						"    Quantization transform function. Extracts and preprocess input data from dataloader item for quantization.\n",
						"    Parameters:\n",
						"       data_item: Dict with data item produced by DataLoader during iteration\n",
						"    Returns:\n",
						"        input_tensor: Input data for quantization\n",
						"    \"\"\"\n",
						"    input_tensor = det_validator.preprocess(data_item)['img'].numpy()\n",
						"    return input_tensor\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    quantization_dataset = nncf.Dataset(det_data_loader, transform_fn)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "ff8a6372-2097-4308-a4e5-e4651242a8df",
					"source": [
						"### Quantize and Save INT8 model\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The `nncf.quantize` function provides an interface for model quantization. It requires an instance of the OpenVINO Model and quantization dataset. \n",
						"Optionally, some additional parameters for the configuration quantization process (number of samples for quantization, preset, ignored scope, etc.) can be provided. YOLOv10 model contains non-ReLU activation functions, which require asymmetric quantization of activations. To achieve a better result, we will use a `mixed` quantization preset. It provides symmetric quantization of weights and asymmetric quantization of activations.\n",
						"\n",
						">**Note**: Model post-training quantization is time-consuming process. Be patient, it can take several minutes depending on your hardware."
					]
				},
				{
					"cell_type": "code",
					"id": "4c12ae38-239e-4fe1-8296-47567f98538c",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"import shutil\n",
						"\n",
						"if not int8_model_det_path.exists():\n",
						"    quantized_det_model = nncf.quantize(\n",
						"        ov_model,\n",
						"        quantization_dataset,\n",
						"        preset=nncf.QuantizationPreset.MIXED,\n",
						"    )\n",
						"\n",
						"    ov.save_model(quantized_det_model,  int8_model_det_path)\n",
						"    shutil.copy(ov_model_path.parent / \"metadata.yaml\", int8_model_det_path.parent / \"metadata.yaml\")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "41b8a3eb-adb8-40cf-8602-311ff8d75a76",
					"source": [
						"## Run Optimized Model Inference\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The way of usage INT8 quantized model is the same like for model before quantization. Let's check inference result of quantized model on single image"
					]
				},
				{
					"cell_type": "markdown",
					"id": "f438df41-d614-4cca-a746-a173d92dddf5",
					"source": [
						"### Run Optimized Model on AUTO device\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "af61303b-16bd-4206-b926-2c28b3a859a4",
					"source": [
						"%%skip not $to_quantize.value\n",
						"ov_yolo_int8_model = YOLOv10(int8_model_det_path.parent, task=\"detect\")"
					]
				},
				{
					"cell_type": "code",
					"id": "db13b8dd-49a1-4046-8a1a-491eee095a1b",
					"source": [
						"%%skip not $to_quantize.value\n",
						"res = ov_yolo_int8_model(IMAGE_PATH, iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "0423ba31-f0ee-4bfc-a18a-f7b870a9683d",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "4b4f7118-e046-4f87-b403-87053f2b0ac3",
					"source": [
						"### Run Optimized Model Inference on selected device\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "142f22ae-6a8d-4578-b97f-4d9c77123418",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"device"
					]
				},
				{
					"cell_type": "code",
					"id": "2f43767a-1994-4202-a411-4738c92223da",
					"source": [
						"%%skip not $to_quantize.value\n",
						"\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value or \"NPU\" in device.value:\n",
						"    ov_model.reshape({0: [1, 3, 640, 640]})\n",
						"ov_config = {}\n",
						"if \"GPU\" in device.value:\n",
						"    ov_config = {\"GPU_DISABLE_WINOGRAD_CONVOLUTION\": \"YES\"}\n",
						"\n",
						"quantized_det_model = core.read_model(int8_model_det_path)\n",
						"quantized_det_compiled_model = core.compile_model(quantized_det_model, device.value, ov_config)\n",
						"\n",
						"ov_yolo_int8_model.predictor.model.ov_compiled_model = quantized_det_compiled_model\n",
						"\n",
						"res = ov_yolo_int8_model(IMAGE_PATH,  iou=0.45, conf=0.2)"
					]
				},
				{
					"cell_type": "code",
					"id": "51d64c7b-7595-41ca-87dd-4c85308f84d6",
					"source": [
						"Image.fromarray(res[0].plot()[:, :, ::-1])"
					]
				},
				{
					"cell_type": "markdown",
					"id": "c7b70f75-a706-4f24-a487-88a3ad645dd5",
					"source": [
						"## Compare the Original and Quantized Models\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "markdown",
					"id": "d6561246-a9e0-4cdf-8207-7e2f919d8582",
					"source": [
						"### Model size\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "43d84d15-75be-4430-a58b-61cfd8e08dd0",
					"source": [
						"ov_model_weights = ov_model_path.with_suffix(\".bin\")\n",
						"print(f\"Size of FP16 model is {ov_model_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
						"if int8_model_det_path.exists():\n",
						"    ov_int8_weights = int8_model_det_path.with_suffix(\".bin\")\n",
						"    print(f\"Size of model with INT8 compressed weights is {ov_int8_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
						"    print(f\"Compression rate for INT8 model: {ov_model_weights.stat().st_size / ov_int8_weights.stat().st_size:.3f}\")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "f72b3d84-91f3-493c-8f4e-46598ffbd6d1",
					"source": [
						"### Performance\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"### FP16 model performance\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "909ea508-e007-4313-ab0f-bfa0e3906176",
					"source": [
						"!benchmark_app -m $ov_model_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "markdown",
					"id": "9852a8ef-4c45-43dd-a9ef-53ca6f13c99e",
					"source": [
						"### Int8 model performance\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "8254626d-6c4f-4ca3-936f-a7c47e402e03",
					"source": [
						"if int8_model_det_path.exists():\n",
						"    !benchmark_app -m $int8_model_det_path -d $device.value -api async -shape \"[1,3,640,640]\" -t 15"
					]
				},
				{
					"cell_type": "markdown",
					"id": "12d35bc1-8101-45ea-8bd7-53720b98f5e5",
					"source": [
						"## Live demo\n",
						"[back to top ⬆️](#Table-of-contents:)\n",
						"\n",
						"\n",
						"The following code runs model inference on a video:"
					]
				},
				{
					"cell_type": "code",
					"id": "246bd379-b1d5-4eb2-928f-97c7c2455a92",
					"source": [
						"import collections\n",
						"import time\n",
						"from IPython import display\n",
						"import cv2\n",
						"import numpy as np\n",
						"\n",
						"\n",
						"# Main processing function to run object detection.\n",
						"def run_object_detection(\n",
						"    source=0,\n",
						"    flip=False,\n",
						"    use_popup=False,\n",
						"    skip_first_frames=0,\n",
						"    det_model=ov_yolo_int8_model,\n",
						"    device=device.value,\n",
						"):\n",
						"    player = None\n",
						"    try:\n",
						"        # Create a video player to play with target fps.\n",
						"        player = VideoPlayer(source=source, flip=flip, fps=30, skip_first_frames=skip_first_frames)\n",
						"        # Start capturing.\n",
						"        player.start()\n",
						"        if use_popup:\n",
						"            title = \"Press ESC to Exit\"\n",
						"            cv2.namedWindow(winname=title, flags=cv2.WINDOW_GUI_NORMAL | cv2.WINDOW_AUTOSIZE)\n",
						"\n",
						"        processing_times = collections.deque()\n",
						"        while True:\n",
						"            # Grab the frame.\n",
						"            frame = player.next()\n",
						"            if frame is None:\n",
						"                print(\"Source ended\")\n",
						"                break\n",
						"            # If the frame is larger than full HD, reduce size to improve the performance.\n",
						"            scale = 1280 / max(frame.shape)\n",
						"            if scale < 1:\n",
						"                frame = cv2.resize(\n",
						"                    src=frame,\n",
						"                    dsize=None,\n",
						"                    fx=scale,\n",
						"                    fy=scale,\n",
						"                    interpolation=cv2.INTER_AREA,\n",
						"                )\n",
						"            # Get the results.\n",
						"            input_image = np.array(frame)\n",
						"\n",
						"            start_time = time.time()\n",
						"            detections = det_model(input_image, iou=0.45, conf=0.2, verbose=False)\n",
						"            stop_time = time.time()\n",
						"            frame = detections[0].plot()\n",
						"\n",
						"            processing_times.append(stop_time - start_time)\n",
						"            # Use processing times from last 200 frames.\n",
						"            if len(processing_times) > 200:\n",
						"                processing_times.popleft()\n",
						"\n",
						"            _, f_width = frame.shape[:2]\n",
						"            # Mean processing time [ms].\n",
						"            processing_time = np.mean(processing_times) * 1000\n",
						"            fps = 1000 / processing_time\n",
						"            cv2.putText(\n",
						"                img=frame,\n",
						"                text=f\"Inference time: {processing_time:.1f}ms ({fps:.1f} FPS)\",\n",
						"                org=(20, 40),\n",
						"                fontFace=cv2.FONT_HERSHEY_COMPLEX,\n",
						"                fontScale=f_width / 1000,\n",
						"                color=(0, 0, 255),\n",
						"                thickness=1,\n",
						"                lineType=cv2.LINE_AA,\n",
						"            )\n",
						"            # Use this workaround if there is flickering.\n",
						"            if use_popup:\n",
						"                cv2.imshow(winname=title, mat=frame)\n",
						"                key = cv2.waitKey(1)\n",
						"                # escape = 27\n",
						"                if key == 27:\n",
						"                    break\n",
						"            else:\n",
						"                # Encode numpy array to jpg.\n",
						"                _, encoded_img = cv2.imencode(ext=\".jpg\", img=frame, params=[cv2.IMWRITE_JPEG_QUALITY, 100])\n",
						"                # Create an IPython image.\n",
						"                i = display.Image(data=encoded_img)\n",
						"                # Display the image in this notebook.\n",
						"                display.clear_output(wait=True)\n",
						"                display.display(i)\n",
						"    # ctrl-c\n",
						"    except KeyboardInterrupt:\n",
						"        print(\"Interrupted\")\n",
						"    # any different error\n",
						"    except RuntimeError as e:\n",
						"        print(e)\n",
						"    finally:\n",
						"        if player is not None:\n",
						"            # Stop capturing.\n",
						"            player.stop()\n",
						"        if use_popup:\n",
						"            cv2.destroyAllWindows()"
					]
				},
				{
					"cell_type": "code",
					"id": "8db930f4-fc41-4635-9248-deab2b6cfcb8",
					"source": [
						"use_int8 = widgets.Checkbox(\n",
						"    value=ov_yolo_int8_model is not None,\n",
						"    description=\"Use int8 model\",\n",
						"    disabled=ov_yolo_int8_model is None,\n",
						")\n",
						"\n",
						"use_int8"
					]
				},
				{
					"cell_type": "code",
					"id": "59e37135-1bdc-40ff-8789-b5b4491cab84",
					"source": [
						"WEBCAM_INFERENCE = False\n",
						"\n",
						"if WEBCAM_INFERENCE:\n",
						"    VIDEO_SOURCE = 0  # Webcam\n",
						"else:\n",
						"    download_file(\n",
						"        \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/people.mp4\",\n",
						"        directory=\"data\",\n",
						"    )\n",
						"    VIDEO_SOURCE = \"data/people.mp4\""
					]
				},
				{
					"cell_type": "code",
					"id": "88c6f47c-1c47-451d-89ba-0a29760f5ec4",
					"source": [
						"run_object_detection(\n",
						"    det_model=ov_yolo_model if not use_int8.value else ov_yolo_int8_model,\n",
						"    source=VIDEO_SOURCE,\n",
						"    flip=True,\n",
						"    use_popup=False,\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"id": "9f612d80-68e2-45f8-aa62-9f3a05ca9de8",
					"source": [
						"### Gradio Interactive Demo\n",
						"[back to top ⬆️](#Table-of-contents:)"
					]
				},
				{
					"cell_type": "code",
					"id": "2e5a6249",
					"source": [
						"def yolov10_inference(image, int8, conf_threshold, iou_threshold):\n",
						"    model = ov_yolo_model if not int8 else ov_yolo_int8_model\n",
						"    results = model(source=image, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]\n",
						"    annotated_image = Image.fromarray(results.plot())\n",
						"\n",
						"    return annotated_image"
					]
				},
				{
					"cell_type": "code",
					"id": "ffdbe248-7086-4804-b4b9-c3dd454bf80c",
					"source": [
						"if not Path(\"gradio_helper.py\").exists():\n",
						"    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/yolov10-optimization/gradio_helper.py\")\n",
						"    open(\"gradio_helper.py\", \"w\").write(r.text)\n",
						"\n",
						"from gradio_helper import make_demo\n",
						"\n",
						"demo = make_demo(fn=yolov10_inference, quantized=ov_yolo_int8_model is not None)\n",
						"\n",
						"try:\n",
						"    demo.launch(debug=True)\n",
						"except Exception:\n",
						"    demo.launch(debug=True, share=True)\n",
						"# If you are launching remotely, specify server_name and server_port\n",
						"# EXAMPLE: `demo.launch(server_name='your server name', server_port='server port in int')`\n",
						"# To learn more please refer to the Gradio docs: https://gradio.app/docs/"
					]
				},
				{
					"cell_type": "code",
					"id": "4b123c5e",
					"source": [
						"# please uncomment and run this cell for stopping gradio interface\n",
						"# demo.close()"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 18 },
			{ modified: 19, original: 19 },
			{ modified: 20, original: 20 },
			{ modified: 21, original: 21 },
			{ modified: 22, original: 22 },
			{ modified: 23, original: 23 },
			{ modified: 24, original: 24 },
			{ modified: 25, original: 25 },
			{ modified: 26, original: 26 },
			{ modified: 27, original: 27 },
			{ modified: 28, original: 28 },
			{ modified: 29, original: 29 },
			{ modified: 30, original: 30 },
			{ modified: 31, original: 31 },
			{ modified: 32, original: 32 },
			{ modified: 33, original: 33 },
			{ modified: 34, original: 34 },
			{ modified: 35, original: 35 },
			{ modified: 36, original: 36 },
			{ modified: 37, original: 37 },
			{ modified: 38, original: 38 },
			{ modified: 39, original: 39 },
			{ modified: 40, original: 40 },
			{ modified: 41, original: 41 },
			{ modified: 42, original: 42 },
			{ modified: 43, original: 43 },
			{ modified: 44, original: 44 },
			{ modified: 45, original: 45 },
			{ modified: 46, original: 46 },
			{ modified: 47, original: 47 },
			{ modified: 48, original: 48 },
			{ modified: 49, original: 49 },
			{ modified: 50, original: 50 },
			{ modified: 51, original: 51 },
			{ modified: 52, original: 52 },
			{ modified: 53, original: -1 },
			{ modified: 54, original: -1 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 52, originalLength: 1, modifiedStart: 52, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 53, originalLength: 0, modifiedStart: 53, modifiedLength: 2 } satisfies IDiffChange,
		]);

	});
	test('Misalignment due to 2 deleted cells', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "F0Wut2G5kz2Z"
					},
					"source": [
						"[![Roboflow Notebooks](https://ik.imagekit.io/roboflow/notebooks/template/bannertest2-2.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672932710194)](https://github.com/roboflow/notebooks)\n",
						"\n",
						"# Fast Segment Anything Model (FastSAM)\n",
						"\n",
						"---\n",
						"\n",
						"[![GitHub](https://badges.aleen42.com/src/github.svg)](https://github.com/CASIA-IVA-Lab/FastSAM)\n",
						"[![arXiv](https://img.shields.io/badge/arXiv-2306.12156-b31b1b.svg)](https://arxiv.org/pdf/2306.12156.pdf)\n",
						"[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/yHNPyqazYYU)\n",
						"[![Roboflow](https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg)](https://blog.roboflow.com/how-to-use-fastsam)\n",
						"\n",
						"The Fast Segment Anything Model (FastSAM) is a CNN Segment Anything Model trained by only 2% of the SA-1B dataset published by SAM authors.\n",
						"\n",
						"![fast-segment-anything-model-figure-1](https://media.roboflow.com/notebooks/examples/fast-sam-figure-1.png)\n",
						"\n",
						"The FastSAM authors claim it achieves comparable performance to the SAM method at 50 times the speed.\n",
						"\n",
						"![fast-segment-anything-model-figure-2](https://media.roboflow.com/notebooks/examples/fast-sam-figure-2.png)\n",
						"\n",
						"## Complementary Materials Covering SAM\n",
						"\n",
						"---\n",
						"\n",
						"[![GitHub](https://badges.aleen42.com/src/github.svg)](https://github.com/facebookresearch/segment-anything)\n",
						"[![arXiv](https://img.shields.io/badge/arXiv-2304.02643-b31b1b.svg)](https://arxiv.org/abs/2304.02643)\n",
						"[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-segment-anything-with-sam.ipynb)\n",
						"[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/D-D6ZmadzPE)\n",
						"[![Roboflow](https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg)](https://blog.roboflow.com/how-to-use-segment-anything-model-sam)\n",
						"\n",
						"## Steps in this Tutorial\n",
						"\n",
						"In this tutorial, we are going to cover:\n",
						"\n",
						"- Before you start\n",
						"- Install FastSAM, SAM, and other dependencies\n",
						"- Download FastSAM and SAM weights\n",
						"- Download example data\n",
						"- Imports\n",
						"- Load FastSAM\n",
						"- FastSAM inference\n",
						"- FastSAM box prompt inference\n",
						"- FastSAM point prompt inference\n",
						"- FastSAM text prompt inference\n",
						"- SAM vs FastSAM\n",
						"- Roboflow benchmark dataset\n",
						"\n",
						"## 🔥 Let's begin!"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "W4u6KjiVdNZM"
					},
					"source": [
						"## Before you start\n",
						"\n",
						"Let's make sure that we have access to GPU. We can use `nvidia-smi` command to do that. In case of any problems navigate to `Edit` -> `Notebook settings` -> `Hardware accelerator`, set it to `GPU`, and then click `Save`."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "Frcrk09FhJeV",
						"outputId": "683a378e-4a8f-4239-901b-037279e92e2b"
					},
					"outputs": [],
					"source": [
						"!nvidia-smi"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "PUACTM2zdmJI"
					},
					"source": [
						"**NOTE:** To make it easier for us to manage datasets, images and models we create a `HOME` constant."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "PRtHQpTNdnm7",
						"outputId": "2824c70d-dcf9-4006-b32d-adc6797fd708"
					},
					"outputs": [],
					"source": [
						"import os\n",
						"HOME = os.getcwd()\n",
						"print(\"HOME:\", HOME)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "YN3DPGZSn57p"
					},
					"source": [
						"## Install FastSAM, SAM, and other dependencies"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "Vr3kcq1KfFir",
						"outputId": "234e5336-ba28-4114-d9e7-efd7f2752ca8"
					},
					"outputs": [],
					"source": [
						"%cd {HOME}\n",
						"\n",
						"# install FastSAM\n",
						"!git clone https://github.com/CASIA-IVA-Lab/FastSAM.git\n",
						"!pip -q install -r FastSAM/requirements.txt\n",
						"# install CLIP\n",
						"!pip -q install git+https://github.com/openai/CLIP.git\n",
						"# install SAM\n",
						"!pip -q install git+https://github.com/facebookresearch/segment-anything.git\n",
						"# install other dependencies\n",
						"!pip -q install roboflow supervision jupyter_bbox_widget"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "BMmgSfDWfFis"
					},
					"source": [
						"## Download FastSAM and SAM weights"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "ADpOBElSfFis",
						"outputId": "b3d89d4c-e2ad-4f5e-a910-b52695263151"
					},
					"outputs": [],
					"source": [
						"!mkdir -p {HOME}/weights\n",
						"!wget -P {HOME}/weights -q https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt\n",
						"!wget -P {HOME}/weights -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\n",
						"!ls -lh {HOME}/weights"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "W6w1ck4z72sn"
					},
					"outputs": [],
					"source": [
						"FAST_SAM_CHECKPOINT_PATH = f\"{HOME}/weights/FastSAM.pt\"\n",
						"SAM_SAM_CHECKPOINT_PATH = f\"{HOME}/weights/sam_vit_h_4b8939.pth\""
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "t1ef2h1R9WMb"
					},
					"source": [
						"## Download example data\n",
						"\n",
						"**NONE:** Let's download few example images. Feel free to use your images or videos."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "CL4gVrnl9dmR",
						"outputId": "db8f78d2-eaf4-41de-9d1f-e5fd6b588a87"
					},
					"outputs": [],
					"source": [
						"!mkdir -p {HOME}/data\n",
						"!wget -P {HOME}/data -q https://media.roboflow.com/notebooks/examples/dog.jpeg\n",
						"!wget -P {HOME}/data -q https://media.roboflow.com/notebooks/examples/robot.jpeg\n",
						"!ls -lh {HOME}/data"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "J_a_icikYiVN"
					},
					"source": [
						"## Imports\n",
						"\n",
						"**NOTE:** `FastSAM` code is not distributed via `pip` not it is packaged. Make sure to run code below from `{HOME}/FastSAM` directory. ⚠️"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "u_NVTXOwYppS",
						"outputId": "b86184de-5827-4b9c-eb5c-959539abe699"
					},
					"outputs": [],
					"source": [
						"%cd {HOME}/FastSAM\n",
						"\n",
						"import os\n",
						"import cv2\n",
						"import torch\n",
						"import roboflow\n",
						"import base64\n",
						"\n",
						"import supervision as sv\n",
						"import numpy as np\n",
						"\n",
						"from roboflow import Roboflow\n",
						"from fastsam import FastSAM, FastSAMPrompt\n",
						"from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "a_Iizsy07pb7"
					},
					"source": [
						"## Load FastSAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "ao8tVDVq7ebG"
					},
					"outputs": [],
					"source": [
						"DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
						"print(f\"DEVICE = {DEVICE}\")\n",
						"fast_sam = FastSAM(FAST_SAM_CHECKPOINT_PATH)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "SvsRFtw_--Nn"
					},
					"source": [
						"## FastSAM inference\n",
						"\n",
						"* `retina_masks=True` determines whether the model uses retina masks for generating segmentation masks.\n",
						"* `imgsz=1024` sets the input image size to 1024x1024 pixels for processing by the model.\n",
						"* `conf=0.4` sets the minimum confidence threshold for object detection.\n",
						"* `iou=0.9` sets the minimum intersection over union threshold for non-maximum suppression to filter out duplicate detections."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "L49P5e-0Iaea"
					},
					"outputs": [],
					"source": [
						"IMAGE_PATH = f\"{HOME}/data/dog.jpeg\""
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"os.makedirs(f\"{HOME}/output\", exist_ok=True)\n",
						"output_image_path = f\"{HOME}/output/output-{os.path.basename(IMAGE_PATH)}\""
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "KV4XzCg1_fAS",
						"outputId": "1df75d90-cf73-4d78-feed-01911a81f6e4"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.4,\n",
						"    iou=0.9)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()\n",
						"prompt_process.plot(annotations=masks, output=f\"{HOME}/output\")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "uPrxqbaBOBMw"
					},
					"source": [
						"**NOTE:** `prompt_process.everything_prompt` returns `torch.Tensor` when DEVICE = 'cuda:0'. `prompt_process.everything_prompt` returns `numpy.ndarray` when DEVICE = 'cpu'."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"# Convert masks to boolean (True/False)\n",
						"def masks_to_bool(masks):\n",
						"    if type(masks) == np.ndarray:\n",
						"        masks = masks.astype(bool)\n",
						"    else:\n",
						"        masks = masks.cpu().numpy().astype(bool)\n",
						"    return masks"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "PUJkojNiKBQa"
					},
					"outputs": [],
					"source": [
						"def annotate_image(image_path: str, masks: np.ndarray) -> np.ndarray:\n",
						"    image = cv2.imread(image_path)\n",
						"\n",
						"    xyxy = sv.mask_to_xyxy(masks=masks)\n",
						"    detections = sv.Detections(xyxy=xyxy, mask=masks)\n",
						"\n",
						"    mask_annotator = sv.MaskAnnotator(color_lookup = sv.ColorLookup.INDEX)\n",
						"    return mask_annotator.annotate(scene=image.copy(), detections=detections)"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "wwtmiH6pLA8N",
						"outputId": "17948850-f051-4a1f-b2b8-d01c2a1f4f20"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "jYkdczKoRUVB"
					},
					"source": [
						"**NOTE:** The quality of the generated masks is quite poor. Let's see if we can improve it by manipulating the parameters."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "dDFGi6EhO2ul",
						"outputId": "f5a03357-eb37-446e-dce7-0659e4c2a368"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "OiBPjQW-O59c",
						"outputId": "afb011fc-933e-4d69-b4dd-e5e2713f407b"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "GoymBQ16KJbv"
					},
					"source": [
						"## FastSAM box prompt inference"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "w1_T5SXpMNku"
					},
					"source": [
						"### Draw Box"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "04-1G1WrMJUm"
					},
					"outputs": [],
					"source": [
						"# helper function that loads an image before adding it to the widget\n",
						"\n",
						"def encode_image(filepath):\n",
						"    with open(filepath, 'rb') as f:\n",
						"        image_bytes = f.read()\n",
						"    encoded = str(base64.b64encode(image_bytes), 'utf-8')\n",
						"    return \"data:image/jpg;base64,\"+encoded"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "xrxd-Ug9MSWG"
					},
					"source": [
						"**NOTE:** Execute cell below and use your mouse to draw bounding box on the image 👇"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 1000,
							"referenced_widgets": [
								"a747e095e21448f29d8697fb6b1a7a74",
								"6401ca8c2f574ba5ab62f0f06b9e0d06"
							]
						},
						"id": "nLA7xmDIMTfx",
						"outputId": "06b51195-d54f-44ed-9236-258e94d624c5"
					},
					"outputs": [],
					"source": [
						"IS_COLAB = True\n",
						"\n",
						"if IS_COLAB:\n",
						"    from google.colab import output\n",
						"    output.enable_custom_widget_manager()\n",
						"\n",
						"from jupyter_bbox_widget import BBoxWidget\n",
						"\n",
						"widget = BBoxWidget()\n",
						"widget.image = encode_image(IMAGE_PATH)\n",
						"widget"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "td6tYT4INOyD",
						"outputId": "73a4e886-dcb8-447c-d3f7-1b3a6496e44d"
					},
					"outputs": [],
					"source": [
						"widget.bboxes"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "XfUvm9bwNUAe"
					},
					"source": [
						"### Generate mask with FastSAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "bI6kPDSROBI-"
					},
					"outputs": [],
					"source": [
						"# default_box is going to be used if you will not draw any box on image above\n",
						"default_box = {'x': 68, 'y': 247, 'width': 555, 'height': 678, 'label': ''}\n",
						"\n",
						"box = widget.bboxes[0] if widget.bboxes else default_box\n",
						"box = [\n",
						"    box['x'],\n",
						"    box['y'],\n",
						"    box['x'] + box['width'],\n",
						"    box['y'] + box['height']\n",
						"]"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "LF879hDqO5iB",
						"outputId": "4e7498c4-4007-4614-d46c-0e35608a04a2"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.box_prompt(bbox=box)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "yt6i2sNJPEdk"
					},
					"source": [
						"**NOTE:** `prompt_process.box_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "EJThTct_PA68",
						"outputId": "f6eddd24-f603-4702-8325-6cff7925fe0f"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "vR-THJthhWH5"
					},
					"source": [
						"## FastSAM point prompt inference"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "xyfpOPPwhpMy"
					},
					"source": [
						"### Select point"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "OC7d3pZ3hop9"
					},
					"outputs": [],
					"source": [
						"point = [405, 505]"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "rE5tV4e3h6nG"
					},
					"source": [
						"### Generate mask with FastSAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "31FIby6Zh-Wa",
						"outputId": "a12c4d1c-837b-4b1b-9240-664797442413"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.point_prompt(points=[point], pointlabel=[1])"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "d7tpVo8CiM3s"
					},
					"source": [
						"**NOTE:** `prompt_process.point_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "lMVpGxShiPky",
						"outputId": "224d84ee-e664-400b-c0df-49507d229f62"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "HOC7v6ATHEr6"
					},
					"source": [
						"## FastSAM text prompt inference"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "4CzLLtXOHeGk",
						"outputId": "cecc50d8-e2be-4b33-850a-a2442e3e30b6"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.text_prompt(text='cap')"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "KYW79zD0OKsv"
					},
					"source": [
						"**NOTE:** `prompt_process.text_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "zODY5xU1LMCb",
						"outputId": "a01540a7-deb3-4555-916c-2ac59928e0f3"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "XFOUWUMZEk7m"
					},
					"source": [
						"##  SAM vs FastSAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "7PKBZ7Hhahhw"
					},
					"outputs": [],
					"source": [
						"IMAGE_PATH = f\"{HOME}/data/robot.jpeg\""
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "c48hRY7dVJeL"
					},
					"source": [
						"### Load SAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "PdKpFCCmaETG"
					},
					"outputs": [],
					"source": [
						"MODEL_TYPE = \"vit_h\"\n",
						"\n",
						"sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_SAM_CHECKPOINT_PATH).to(device=DEVICE)\n",
						"mask_generator = SamAutomaticMaskGenerator(sam)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "V5bPGvtBaewO"
					},
					"source": [
						"### SAM inference"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "0HgWsk1yalow"
					},
					"outputs": [],
					"source": [
						"image_bgr = cv2.imread(IMAGE_PATH)\n",
						"image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
						"\n",
						"sam_result = mask_generator.generate(image_rgb)\n",
						"sam_detections = sv.Detections.from_sam(sam_result=sam_result)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "t03RrSzMasN6"
					},
					"source": [
						"### FastSAM inference"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "xuFGLw8MavQI",
						"outputId": "b0118973-1e29-419d-b7c0-2933b97d7ce5"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()\n",
						"masks = masks_to_bool(masks)\n",
						"xyxy = sv.mask_to_xyxy(masks=masks)\n",
						"fast_sam_detections = sv.Detections(xyxy=xyxy, mask=masks)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "IYn8xKySbBnc"
					},
					"source": [
						"### FastSAM vs. SAM"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 321
						},
						"id": "c2Qj7sLjbDvo",
						"outputId": "e85e3ed3-8b90-409d-cf95-877a6a9b2034"
					},
					"outputs": [],
					"source": [
						"mask_annotator = sv.MaskAnnotator(color_lookup = sv.ColorLookup.INDEX)\n",
						"\n",
						"sam_result = mask_annotator.annotate(scene=image_bgr.copy(), detections=sam_detections)\n",
						"fast_sam_result = mask_annotator.annotate(scene=image_bgr.copy(), detections=fast_sam_detections)\n",
						"\n",
						"sv.plot_images_grid(\n",
						"    images=[sam_result, fast_sam_result],\n",
						"    grid_size=(1, 2),\n",
						"    titles=['SAM', 'FastSAM']\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "MfU4NQiGWQMC"
					},
					"source": [
						"**NOTE:** There is a lot going on in our test image. Let's plot our masks against a blank background."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 321
						},
						"id": "PIrywlF5QRPz",
						"outputId": "8e7c84ae-7869-4523-847a-0678fd1aea53"
					},
					"outputs": [],
					"source": [
						"mask_annotator = sv.MaskAnnotator(color_lookup = sv.ColorLookup.INDEX)\n",
						"\n",
						"sam_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=sam_detections)\n",
						"fast_sam_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=fast_sam_detections)\n",
						"\n",
						"sv.plot_images_grid(\n",
						"    images=[sam_result, fast_sam_result],\n",
						"    grid_size=(1, 2),\n",
						"    titles=['SAM', 'FastSAM']\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "51Fs3sDGQmIt"
					},
					"source": [
						"### Mask diff\n",
						"\n",
						"**NOTE:** Let's try to understand which masks generated by SAM were not generated by FastSAM."
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "O5bIUAhYQrqZ"
					},
					"outputs": [],
					"source": [
						"def compute_iou(mask1, mask2):\n",
						"    intersection = np.logical_and(mask1, mask2)\n",
						"    union = np.logical_or(mask1, mask2)\n",
						"    return np.sum(intersection) / np.sum(union)\n",
						"\n",
						"def filter_masks(mask_array1, mask_array2, threshold):\n",
						"    return np.array([mask1 for mask1 in mask_array1 if max(compute_iou(mask1, mask2) for mask2 in mask_array2) < threshold])"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 420
						},
						"id": "1GM4bL9sRbOm",
						"outputId": "421fb4f5-7a42-461b-ba1e-7cd46d2d7bd7"
					},
					"outputs": [],
					"source": [
						"diff_masks = filter_masks(sam_detections.mask, fast_sam_detections.mask, 0.5)\n",
						"diff_xyxy = sv.mask_to_xyxy(masks=diff_masks)\n",
						"diff_detections = sv.Detections(xyxy=diff_xyxy, mask=diff_masks)\n",
						"\n",
						"mask_annotator = sv.MaskAnnotator(color_lookup = sv.ColorLookup.INDEX)\n",
						"diff_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=diff_detections)\n",
						"sv.plot_image(image=diff_result, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "0lPZAZZulhF8"
					},
					"source": [
						"## Roboflow Benchmark Dataset"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "LWN6XwPuSZBY",
						"outputId": "86b18c0e-f6d6-486b-eee6-d004f13f531e"
					},
					"outputs": [],
					"source": [
						"%cd {HOME}\n",
						"\n",
						"roboflow.login()\n",
						"rf = Roboflow()\n",
						"\n",
						"project = rf.workspace(\"roboticfish\").project(\"underwater_object_detection\")\n",
						"dataset = project.version(10).download(\"yolov8\")"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"id": "T6-PaxWXTZ69"
					},
					"outputs": [],
					"source": [
						"IMAGE_PATH = os.path.join(dataset.location, 'train', 'images', 'IMG_2311_jpeg_jpg.rf.09ae6820eaff21dc838b1f9b6b20342b.jpg')"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "p-gY5jf1Tife",
						"outputId": "630d35d8-26f4-4342-8d82-c7b61be26fe8"
					},
					"outputs": [],
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.text_prompt(text='Penguin')"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "3tBIH0yubHsW",
						"outputId": "fc4d0a6a-53c1-4a2b-91f0-75c8f78e10fd"
					},
					"outputs": [],
					"source": [
						"masks = masks_to_bool(masks)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "gP9gEw9Tp8GJ"
					},
					"source": [
						"## 🏆 Congratulations\n",
						"\n",
						"### Learning Resources\n",
						"\n",
						"Roboflow has produced many resources that you may find interesting as you advance your knowledge of computer vision:\n",
						"\n",
						"- [Roboflow Notebooks](https://github.com/roboflow/notebooks): A repository of over 20 notebooks that walk through how to train custom models with a range of model types, from YOLOv7 to SegFormer.\n",
						"- [Roboflow YouTube](https://www.youtube.com/c/Roboflow): Our library of videos featuring deep dives into the latest in computer vision, detailed tutorials that accompany our notebooks, and more.\n",
						"- [Roboflow Discuss](https://discuss.roboflow.com/): Have a question about how to do something on Roboflow? Ask your question on our discussion forum.\n",
						"- [Roboflow Models](https://roboflow.com): Learn about state-of-the-art models and their performance. Find links and tutorials to guide your learning.\n",
						"\n",
						"### Convert data formats\n",
						"\n",
						"Roboflow provides free utilities to convert data between dozens of popular computer vision formats. Check out [Roboflow Formats](https://roboflow.com/formats) to find tutorials on how to convert data between formats in a few clicks.\n",
						"\n",
						"### Connect computer vision to your project logic\n",
						"\n",
						"[Roboflow Templates](https://roboflow.com/templates) is a public gallery of code snippets that you can use to connect computer vision to your project logic. Code snippets range from sending emails after inference to measuring object distance between detections."
					]
				}
			]
				.map(fromJupyterCell)
			,
			[
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "F0Wut2G5kz2Z"
					},
					"source": [
						"[![Roboflow Notebooks](https://ik.imagekit.io/roboflow/notebooks/template/bannertest2-2.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672932710194)](https://github.com/roboflow/notebooks)\n",
						"\n",
						"# Fast Segment Anything Model (FastSAM)\n",
						"\n",
						"---\n",
						"\n",
						"[![GitHub](https://badges.aleen42.com/src/github.svg)](https://github.com/CASIA-IVA-Lab/FastSAM)\n",
						"[![arXiv](https://img.shields.io/badge/arXiv-2306.12156-b31b1b.svg)](https://arxiv.org/pdf/2306.12156.pdf)\n",
						"[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/yHNPyqazYYU)\n",
						"[![Roboflow](https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg)](https://blog.roboflow.com/how-to-use-fastsam)\n",
						"\n",
						"The Fast Segment Anything Model (FastSAM) is a CNN Segment Anything Model trained by only 2% of the SA-1B dataset published by SAM authors.\n",
						"\n",
						"![fast-segment-anything-model-figure-1](https://media.roboflow.com/notebooks/examples/fast-sam-figure-1.png)\n",
						"\n",
						"The FastSAM authors claim it achieves comparable performance to the SAM method at 50 times the speed.\n",
						"\n",
						"![fast-segment-anything-model-figure-2](https://media.roboflow.com/notebooks/examples/fast-sam-figure-2.png)\n",
						"\n",
						"## Complementary Materials Covering SAM\n",
						"\n",
						"---\n",
						"\n",
						"[![GitHub](https://badges.aleen42.com/src/github.svg)](https://github.com/facebookresearch/segment-anything)\n",
						"[![arXiv](https://img.shields.io/badge/arXiv-2304.02643-b31b1b.svg)](https://arxiv.org/abs/2304.02643)\n",
						"[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-segment-anything-with-sam.ipynb)\n",
						"[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://youtu.be/D-D6ZmadzPE)\n",
						"[![Roboflow](https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg)](https://blog.roboflow.com/how-to-use-segment-anything-model-sam)\n",
						"\n",
						"## Steps in this Tutorial\n",
						"\n",
						"In this tutorial, we are going to cover:\n",
						"\n",
						"- Before you start\n",
						"- Install FastSAM, SAM, and other dependencies\n",
						"- Download FastSAM and SAM weights\n",
						"- Download example data\n",
						"- Imports\n",
						"- Load FastSAM\n",
						"- FastSAM inference\n",
						"- FastSAM box prompt inference\n",
						"- FastSAM point prompt inference\n",
						"- FastSAM text prompt inference\n",
						"- SAM vs FastSAM\n",
						"- Roboflow benchmark dataset\n",
						"\n",
						"## 🔥 Let's begin!"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "W4u6KjiVdNZM"
					},
					"source": [
						"## Before you start\n",
						"\n",
						"Let's make sure that we have access to GPU. We can use `nvidia-smi` command to do that. In case of any problems navigate to `Edit` -> `Notebook settings` -> `Hardware accelerator`, set it to `GPU`, and then click `Save`."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "Frcrk09FhJeV",
						"outputId": "683a378e-4a8f-4239-901b-037279e92e2b"
					},
					"source": [
						"!nvidia-smi"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "PUACTM2zdmJI"
					},
					"source": [
						"**NOTE:** To make it easier for us to manage datasets, images and models we create a `HOME` constant."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "PRtHQpTNdnm7",
						"outputId": "2824c70d-dcf9-4006-b32d-adc6797fd708"
					},
					"source": [
						"import os\n",
						"HOME = os.getcwd()\n",
						"print(\"HOME:\", HOME)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "YN3DPGZSn57p"
					},
					"source": [
						"## Install FastSAM, SAM, and other dependencies"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "Vr3kcq1KfFir",
						"outputId": "234e5336-ba28-4114-d9e7-efd7f2752ca8"
					},
					"source": [
						"%cd {HOME}\n",
						"\n",
						"# install FastSAM\n",
						"!git clone https://github.com/CASIA-IVA-Lab/FastSAM.git\n",
						"!pip -q install -r FastSAM/requirements.txt\n",
						"# install CLIP\n",
						"!pip -q install git+https://github.com/openai/CLIP.git\n",
						"# install SAM\n",
						"!pip -q install git+https://github.com/facebookresearch/segment-anything.git\n",
						"# install other dependencies\n",
						"!pip -q install roboflow supervision jupyter_bbox_widget"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "BMmgSfDWfFis"
					},
					"source": [
						"## Download FastSAM and SAM weights"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "ADpOBElSfFis",
						"outputId": "b3d89d4c-e2ad-4f5e-a910-b52695263151"
					},
					"source": [
						"!mkdir -p {HOME}/weights\n",
						"!wget -P {HOME}/weights -q https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt\n",
						"!wget -P {HOME}/weights -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\n",
						"!ls -lh {HOME}/weights"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "W6w1ck4z72sn"
					},
					"source": [
						"FAST_SAM_CHECKPOINT_PATH = f\"{HOME}/weights/FastSAM.pt\"\n",
						"SAM_SAM_CHECKPOINT_PATH = f\"{HOME}/weights/sam_vit_h_4b8939.pth\""
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "t1ef2h1R9WMb"
					},
					"source": [
						"## Download example data\n",
						"\n",
						"**NONE:** Let's download few example images. Feel free to use your images or videos."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "CL4gVrnl9dmR",
						"outputId": "db8f78d2-eaf4-41de-9d1f-e5fd6b588a87"
					},
					"source": [
						"!mkdir -p {HOME}/data\n",
						"!wget -P {HOME}/data -q https://media.roboflow.com/notebooks/examples/dog.jpeg\n",
						"!wget -P {HOME}/data -q https://media.roboflow.com/notebooks/examples/robot.jpeg\n",
						"!ls -lh {HOME}/data"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "J_a_icikYiVN"
					},
					"source": [
						"## Imports\n",
						"\n",
						"**NOTE:** `FastSAM` code is not distributed via `pip` not it is packaged. Make sure to run code balow from `{HOME}/FastSAM` directory. ⚠️"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "u_NVTXOwYppS",
						"outputId": "b86184de-5827-4b9c-eb5c-959539abe699"
					},
					"source": [
						"%cd {HOME}/FastSAM\n",
						"\n",
						"import os\n",
						"import cv2\n",
						"import torch\n",
						"import roboflow\n",
						"import base64\n",
						"\n",
						"import supervision as sv\n",
						"import numpy as np\n",
						"\n",
						"from roboflow import Roboflow\n",
						"from fastsam import FastSAM, FastSAMPrompt\n",
						"from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "a_Iizsy07pb7"
					},
					"source": [
						"## Load FastSAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "ao8tVDVq7ebG"
					},
					"source": [
						"DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
						"\n",
						"fast_sam = FastSAM(FAST_SAM_CHECKPOINT_PATH)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "SvsRFtw_--Nn"
					},
					"source": [
						"## FastSAM inference\n",
						"\n",
						"* `retina_masks=True` determines whether the model uses retina masks for generating segmentation masks.\n",
						"* `imgsz=1024` sets the input image size to 1024x1024 pixels for processing by the model.\n",
						"* `conf=0.4` sets the minimum confidence threshold for object detection.\n",
						"* `iou=0.9` sets the minimum intersection over union threshold for non-maximum suppression to filter out duplicate detections."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "L49P5e-0Iaea"
					},
					"source": [
						"IMAGE_PATH = f\"{HOME}/data/dog.jpeg\""
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "KV4XzCg1_fAS",
						"outputId": "1df75d90-cf73-4d78-feed-01911a81f6e4"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.4,\n",
						"    iou=0.9)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()\n",
						"prompt_process.plot(annotations=masks, output=f\"{HOME}/output\")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "uPrxqbaBOBMw"
					},
					"source": [
						"**NOTE:** `prompt_process.everything_prompt` returns `torch.Tensor`"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "PUJkojNiKBQa"
					},
					"source": [
						"def annotate_image(image_path: str, masks: np.ndarray) -> np.ndarray:\n",
						"    image = cv2.imread(image_path)\n",
						"\n",
						"    xyxy = sv.mask_to_xyxy(masks=masks)\n",
						"    detections = sv.Detections(xyxy=xyxy, mask=masks)\n",
						"\n",
						"    mask_annotator = sv.MaskAnnotator()\n",
						"    return mask_annotator.annotate(scene=image.copy(), detections=detections)"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "wwtmiH6pLA8N",
						"outputId": "17948850-f051-4a1f-b2b8-d01c2a1f4f20"
					},
					"source": [
						"masks = masks.cpu().numpy().astype(bool)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "jYkdczKoRUVB"
					},
					"source": [
						"**NOTE:** The quality of the generated masks is quite poor. Let's see if we can improve it by manipulating the parameters."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "dDFGi6EhO2ul",
						"outputId": "f5a03357-eb37-446e-dce7-0659e4c2a368"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "OiBPjQW-O59c",
						"outputId": "afb011fc-933e-4d69-b4dd-e5e2713f407b"
					},
					"source": [
						"masks = masks.cpu().numpy().astype(bool)\n",
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "GoymBQ16KJbv"
					},
					"source": [
						"## FastSAM box prompt inference"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "w1_T5SXpMNku"
					},
					"source": [
						"### Draw Box"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "04-1G1WrMJUm"
					},
					"source": [
						"# helper function that loads an image before adding it to the widget\n",
						"\n",
						"def encode_image(filepath):\n",
						"    with open(filepath, 'rb') as f:\n",
						"        image_bytes = f.read()\n",
						"    encoded = str(base64.b64encode(image_bytes), 'utf-8')\n",
						"    return \"data:image/jpg;base64,\"+encoded"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "xrxd-Ug9MSWG"
					},
					"source": [
						"**NOTE:** Execute cell below and use your mouse to draw bounding box on the image 👇"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 1000,
							"referenced_widgets": [
								"a747e095e21448f29d8697fb6b1a7a74",
								"6401ca8c2f574ba5ab62f0f06b9e0d06"
							]
						},
						"id": "nLA7xmDIMTfx",
						"outputId": "06b51195-d54f-44ed-9236-258e94d624c5"
					},
					"source": [
						"IS_COLAB = True\n",
						"\n",
						"if IS_COLAB:\n",
						"    from google.colab import output\n",
						"    output.enable_custom_widget_manager()\n",
						"\n",
						"from jupyter_bbox_widget import BBoxWidget\n",
						"\n",
						"widget = BBoxWidget()\n",
						"widget.image = encode_image(IMAGE_PATH)\n",
						"widget"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "td6tYT4INOyD",
						"outputId": "73a4e886-dcb8-447c-d3f7-1b3a6496e44d"
					},
					"source": [
						"widget.bboxes"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "XfUvm9bwNUAe"
					},
					"source": [
						"### Generate mask with FastSAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "bI6kPDSROBI-"
					},
					"source": [
						"# default_box is going to be used if you will not draw any box on image above\n",
						"default_box = {'x': 68, 'y': 247, 'width': 555, 'height': 678, 'label': ''}\n",
						"\n",
						"box = widget.bboxes[0] if widget.bboxes else default_box\n",
						"box = [\n",
						"    box['x'],\n",
						"    box['y'],\n",
						"    box['x'] + box['width'],\n",
						"    box['y'] + box['height']\n",
						"]"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "LF879hDqO5iB",
						"outputId": "4e7498c4-4007-4614-d46c-0e35608a04a2"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.box_prompt(bbox=box)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "yt6i2sNJPEdk"
					},
					"source": [
						"**NOTE:** `prompt_process.box_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "EJThTct_PA68",
						"outputId": "f6eddd24-f603-4702-8325-6cff7925fe0f"
					},
					"source": [
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "vR-THJthhWH5"
					},
					"source": [
						"## FastSAM point prompt inference"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "xyfpOPPwhpMy"
					},
					"source": [
						"### Select point"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "OC7d3pZ3hop9"
					},
					"source": [
						"point = [405, 505]"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "rE5tV4e3h6nG"
					},
					"source": [
						"### Generate mask with FastSAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "31FIby6Zh-Wa",
						"outputId": "a12c4d1c-837b-4b1b-9240-664797442413"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.point_prompt(points=[point], pointlabel=[1])"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "d7tpVo8CiM3s"
					},
					"source": [
						"**NOTE:** `prompt_process.point_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "lMVpGxShiPky",
						"outputId": "224d84ee-e664-400b-c0df-49507d229f62"
					},
					"source": [
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "HOC7v6ATHEr6"
					},
					"source": [
						"## FastSAM text prompt inference"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "4CzLLtXOHeGk",
						"outputId": "cecc50d8-e2be-4b33-850a-a2442e3e30b6"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.text_prompt(text='cap')"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "KYW79zD0OKsv"
					},
					"source": [
						"**NOTE:** `prompt_process.text_prompt` returns `numy.ndarray`"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "zODY5xU1LMCb",
						"outputId": "a01540a7-deb3-4555-916c-2ac59928e0f3"
					},
					"source": [
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "XFOUWUMZEk7m"
					},
					"source": [
						"##  SAM vs FastSAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "7PKBZ7Hhahhw"
					},
					"source": [
						"IMAGE_PATH = f\"{HOME}/data/robot.jpeg\""
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "c48hRY7dVJeL"
					},
					"source": [
						"### Load SAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "PdKpFCCmaETG"
					},
					"source": [
						"MODEL_TYPE = \"vit_h\"\n",
						"\n",
						"sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_SAM_CHECKPOINT_PATH).to(device=DEVICE)\n",
						"mask_generator = SamAutomaticMaskGenerator(sam)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "V5bPGvtBaewO"
					},
					"source": [
						"### SAM inference"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "0HgWsk1yalow"
					},
					"source": [
						"image_bgr = cv2.imread(IMAGE_PATH)\n",
						"image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
						"\n",
						"sam_result = mask_generator.generate(image_rgb)\n",
						"sam_detections = sv.Detections.from_sam(sam_result=sam_result)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "t03RrSzMasN6"
					},
					"source": [
						"### FastSAM inference"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "xuFGLw8MavQI",
						"outputId": "b0118973-1e29-419d-b7c0-2933b97d7ce5"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.everything_prompt()\n",
						"masks = masks.cpu().numpy().astype(bool)\n",
						"xyxy = sv.mask_to_xyxy(masks=masks)\n",
						"fast_sam_detections = sv.Detections(xyxy=xyxy, mask=masks)"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "IYn8xKySbBnc"
					},
					"source": [
						"### FastSAM vs. SAM"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 321
						},
						"id": "c2Qj7sLjbDvo",
						"outputId": "e85e3ed3-8b90-409d-cf95-877a6a9b2034"
					},
					"source": [
						"mask_annotator = sv.MaskAnnotator()\n",
						"\n",
						"sam_result = mask_annotator.annotate(scene=image_bgr.copy(), detections=sam_detections)\n",
						"fast_sam_result = mask_annotator.annotate(scene=image_bgr.copy(), detections=fast_sam_detections)\n",
						"\n",
						"sv.plot_images_grid(\n",
						"    images=[sam_result, fast_sam_result],\n",
						"    grid_size=(1, 2),\n",
						"    titles=['SAM', 'FastSAM']\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "MfU4NQiGWQMC"
					},
					"source": [
						"**NOTE:** There is a lot going on in our test image. Let's plot our masks against a blank background."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 321
						},
						"id": "PIrywlF5QRPz",
						"outputId": "8e7c84ae-7869-4523-847a-0678fd1aea53"
					},
					"source": [
						"mask_annotator = sv.MaskAnnotator()\n",
						"\n",
						"sam_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=sam_detections)\n",
						"fast_sam_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=fast_sam_detections)\n",
						"\n",
						"sv.plot_images_grid(\n",
						"    images=[sam_result, fast_sam_result],\n",
						"    grid_size=(1, 2),\n",
						"    titles=['SAM', 'FastSAM']\n",
						")"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "51Fs3sDGQmIt"
					},
					"source": [
						"### Mask diff\n",
						"\n",
						"**NOTE:** Let's try to understand which masks generated by SAM were not generated by FastSAM."
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "O5bIUAhYQrqZ"
					},
					"source": [
						"def compute_iou(mask1, mask2):\n",
						"    intersection = np.logical_and(mask1, mask2)\n",
						"    union = np.logical_or(mask1, mask2)\n",
						"    return np.sum(intersection) / np.sum(union)\n",
						"\n",
						"def filter_masks(mask_array1, mask_array2, threshold):\n",
						"    return np.array([mask1 for mask1 in mask_array1 if max(compute_iou(mask1, mask2) for mask2 in mask_array2) < threshold])"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 420
						},
						"id": "1GM4bL9sRbOm",
						"outputId": "421fb4f5-7a42-461b-ba1e-7cd46d2d7bd7"
					},
					"source": [
						"diff_masks = filter_masks(sam_detections.mask, fast_sam_detections.mask, 0.5)\n",
						"diff_xyxy = sv.mask_to_xyxy(masks=diff_masks)\n",
						"diff_detections = sv.Detections(xyxy=diff_xyxy, mask=diff_masks)\n",
						"\n",
						"mask_annotator = sv.MaskAnnotator()\n",
						"diff_result = mask_annotator.annotate(scene=np.zeros_like(image_bgr), detections=diff_detections)\n",
						"sv.plot_image(image=diff_result, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "0lPZAZZulhF8"
					},
					"source": [
						"## Roboflow Benchmark Dataset"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "LWN6XwPuSZBY",
						"outputId": "86b18c0e-f6d6-486b-eee6-d004f13f531e"
					},
					"source": [
						"%cd {HOME}\n",
						"\n",
						"roboflow.login()\n",
						"rf = Roboflow()\n",
						"\n",
						"project = rf.workspace(\"roboticfish\").project(\"underwater_object_detection\")\n",
						"dataset = project.version(10).download(\"yolov8\")"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"id": "T6-PaxWXTZ69"
					},
					"source": [
						"IMAGE_PATH = os.path.join(dataset.location, 'train', 'images', 'IMG_2311_jpeg_jpg.rf.09ae6820eaff21dc838b1f9b6b20342b.jpg')"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/"
						},
						"id": "p-gY5jf1Tife",
						"outputId": "630d35d8-26f4-4342-8d82-c7b61be26fe8"
					},
					"source": [
						"results = fast_sam(\n",
						"    source=IMAGE_PATH,\n",
						"    device=DEVICE,\n",
						"    retina_masks=True,\n",
						"    imgsz=1024,\n",
						"    conf=0.5,\n",
						"    iou=0.6)\n",
						"prompt_process = FastSAMPrompt(IMAGE_PATH, results, device=DEVICE)\n",
						"masks = prompt_process.text_prompt(text='Penguin')"
					]
				},
				{
					"cell_type": "code",
					"metadata": {
						"colab": {
							"base_uri": "https://localhost:8080/",
							"height": 653
						},
						"id": "3tBIH0yubHsW",
						"outputId": "fc4d0a6a-53c1-4a2b-91f0-75c8f78e10fd"
					},
					"source": [
						"annotated_image=annotate_image(image_path=IMAGE_PATH, masks=masks)\n",
						"sv.plot_image(image=annotated_image, size=(8, 8))"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {
						"id": "gP9gEw9Tp8GJ"
					},
					"source": [
						"## 🏆 Congratulations\n",
						"\n",
						"### Learning Resources\n",
						"\n",
						"Roboflow has produced many resources that you may find interesting as you advance your knowledge of computer vision:\n",
						"\n",
						"- [Roboflow Notebooks](https://github.com/roboflow/notebooks): A repository of over 20 notebooks that walk through how to train custom models with a range of model types, from YOLOv7 to SegFormer.\n",
						"- [Roboflow YouTube](https://www.youtube.com/c/Roboflow): Our library of videos featuring deep dives into the latest in computer vision, detailed tutorials that accompany our notebooks, and more.\n",
						"- [Roboflow Discuss](https://discuss.roboflow.com/): Have a question about how to do something on Roboflow? Ask your question on our discussion forum.\n",
						"- [Roboflow Models](https://roboflow.com): Learn about state-of-the-art models and their performance. Find links and tutorials to guide your learning.\n",
						"\n",
						"### Convert data formats\n",
						"\n",
						"Roboflow provides free utilities to convert data between dozens of popular computer vision formats. Check out [Roboflow Formats](https://roboflow.com/formats) to find tutorials on how to convert data between formats in a few clicks.\n",
						"\n",
						"### Connect computer vision to your project logic\n",
						"\n",
						"[Roboflow Templates](https://roboflow.com/templates) is a public gallery of code snippets that you can use to connect computer vision to your project logic. Code snippets range from sending emails after inference to measuring object distance between detections."
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: 0 },
			{ modified: 1, original: 1 },
			{ modified: 2, original: 2 },
			{ modified: 3, original: 3 },
			{ modified: 4, original: 4 },
			{ modified: 5, original: 5 },
			{ modified: 6, original: 6 },
			{ modified: 7, original: 7 },
			{ modified: 8, original: 8 },
			{ modified: 9, original: 9 },
			{ modified: 10, original: 10 },
			{ modified: 11, original: 11 },
			{ modified: 12, original: 12 },
			{ modified: 13, original: 13 },
			{ modified: 14, original: 14 },
			{ modified: 15, original: 15 },
			{ modified: 16, original: 16 },
			{ modified: 17, original: 17 },
			{ modified: 18, original: 19 },
			{ modified: 19, original: 20 },
			{ modified: 20, original: 22 },
			{ modified: 21, original: 23 },
			{ modified: 22, original: 24 },
			{ modified: 23, original: 25 },
			{ modified: 24, original: 26 },
			{ modified: 25, original: 27 },
			{ modified: 26, original: 28 },
			{ modified: 27, original: 29 },
			{ modified: 28, original: 30 },
			{ modified: 29, original: 31 },
			{ modified: 30, original: 32 },
			{ modified: 31, original: 33 },
			{ modified: 32, original: 34 },
			{ modified: 33, original: 35 },
			{ modified: 34, original: 36 },
			{ modified: 35, original: 37 },
			{ modified: 36, original: 38 },
			{ modified: 37, original: 39 },
			{ modified: 38, original: 40 },
			{ modified: 39, original: 41 },
			{ modified: 40, original: 42 },
			{ modified: 41, original: 43 },
			{ modified: 42, original: 44 },
			{ modified: 43, original: 45 },
			{ modified: 44, original: 46 },
			{ modified: 45, original: 47 },
			{ modified: 46, original: 48 },
			{ modified: 47, original: 49 },
			{ modified: 48, original: 50 },
			{ modified: 49, original: 51 },
			{ modified: 50, original: 52 },
			{ modified: 51, original: 53 },
			{ modified: 52, original: 54 },
			{ modified: 53, original: 55 },
			{ modified: 54, original: 56 },
			{ modified: 55, original: 57 },
			{ modified: 56, original: 58 },
			{ modified: 57, original: 59 },
			{ modified: 58, original: 60 },
			{ modified: 59, original: 61 },
			{ modified: 60, original: 62 },
			{ modified: 61, original: 63 },
			{ modified: 62, original: 64 },
			{ modified: 63, original: 65 },
			{ modified: 64, original: 66 },
			{ modified: 65, original: 67 },
			{ modified: 66, original: 68 },
			{ modified: 67, original: 69 },
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 12, originalLength: 1, modifiedStart: 12, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 15, originalLength: 1, modifiedStart: 15, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 18, originalLength: 1, modifiedStart: 18, modifiedLength: 0 } satisfies IDiffChange,
			{ originalStart: 20, originalLength: 1, modifiedStart: 19, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 21, originalLength: 1, modifiedStart: 20, modifiedLength: 0 } satisfies IDiffChange,
			{ originalStart: 22, originalLength: 1, modifiedStart: 20, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 23, originalLength: 1, modifiedStart: 21, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 26, originalLength: 1, modifiedStart: 24, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 37, originalLength: 1, modifiedStart: 35, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 44, originalLength: 1, modifiedStart: 42, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 48, originalLength: 1, modifiedStart: 46, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 56, originalLength: 1, modifiedStart: 54, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 58, originalLength: 1, modifiedStart: 56, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 60, originalLength: 1, modifiedStart: 58, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 63, originalLength: 1, modifiedStart: 61, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 68, originalLength: 1, modifiedStart: 66, modifiedLength: 1 } satisfies IDiffChange,
			// { originalStart: 53, originalLength: 0, modifiedStart: 53, modifiedLength: 2 } satisfies IDiffChange,
		]);

	});
	test('Insert a few cells (one on top), modify a few and completely misaligned', async () => {
		const { mapping, diff } = await mapCells(
			[
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import requests\n",
						"import pandas as pd\n",
						"import matplotlib.pyplot as plt\n",
						"from cachetools import cached, TTLCache\n",
						"from dotenv import load_dotenv"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import os\n",
						"\n",
						"GITHUB_API_URL = \"https://api.github.com\"\n",
						"REPO_OWNER = \"microsoft\"\n",
						"ISSUE = 123  # Replace with your milestone number\n",
						"load_dotenv()\n",
						"TOKEN = os.getenv(\"GH\")\n",
						"\n",
						"headers = {\n",
						"    \"Authorization\": f\"Bearer {TOKEN}\",\n",
						"    \"Accept\": \"application/vnd.github+json\"\n",
						"}"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import csv\n",
						"from datetime import datetime\n",
						"import os\n",
						"\n",
						"def append_information(repo, open_issues_count, ):\n",
						"    file_exists = os.path.isfile(filename)\n",
						"    timestamp = datetime.utcnow().isoformat()\n",
						"\n",
						"    # Open the file in append mode\n",
						"    with open(filename, 'a', newline='', encoding='utf-8') as csvfile:\n",
						"        writer = csv.writer(csvfile)\n",
						"\n",
						"        # If the file didn't exist before, write the header row\n",
						"        if not file_exists:\n",
						"            writer.writerow([\n",
						"                \"timestamp\",\n",
						"                \"repo\",\n",
						"                \"open_issues_count\",\n",
						"                \"open_prs_count\",\n",
						"                \"bug_issues_count\",\n",
						"                \"feature_request_issues_count\",\n",
						"                \"invalid_issues_count\"\n",
						"            ])\n"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"def get_repo_issue_stats(owner, repo, token):\n",
						"    # Common GitHub GraphQL endpoint\n",
						"    url = \"https://api.github.com/graphql\"\n",
						"    headers = {\n",
						"        \"Authorization\": f\"Bearer {token}\",\n",
						"        \"Accept\": \"application/vnd.github.v4+json\"\n",
						"    }\n",
						"\n",
						"    # First query: Get various issue/PR counts from the repository object\n",
						"    query_repo = \"\"\"\n",
						"    query ($owner: String!, $name: String!) {\n",
						"      repository(owner: $owner, name: $name) {\n",
						"        issues(states: OPEN) {\n",
						"          totalCount\n",
						"        }\n",
						"        pullRequests(states: OPEN) {\n",
						"          totalCount\n",
						"        }\n",
						"        bugIssues: issues(states: OPEN, labels: [\"bug\"]) {\n",
						"          totalCount\n",
						"        }\n",
						"        featureRequestIssues: issues(states: OPEN, labels: [\"feature-request\"]) {\n",
						"          totalCount\n",
						"        }\n",
						"      }\n",
						"    }\n",
						"    \"\"\"\n"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"for REPO in REPOS:\n",
						"    stats = get_repo_issue_stats(OWNER, REPO, TOKEN)\n",
						"    print(stats)\n",
						"\n",
						"    append_metrics_to_csv(\n",
						"        filename=\"repo_metrics.csv\",\n",
						"        repo=f\"{OWNER}/{REPO}\",\n",
						"        open_issues_count=stats['open_issues_count'],\n",
						"        open_prs_count=stats['open_prs_count'],\n",
						"        bug_issues_count=stats['bug_issues_count'],\n",
						"        feature_request_issues_count=stats['feature_request_issues_count'],\n",
						"        invalid_issues_count=stats['invalid_issues_count']\n",
						"    )"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {},
					"source": [
						"# Plot"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import pandas as pd\n",
						"import plotly.graph_objs as go\n",
						"from plotly.subplots import make_subplots\n",
						"\n",
						"def plot_repo_metrics(repo, csv_filename=\"repo_metrics.csv\"):\n",
						"    # Load the CSV data\n",
						"    df = pd.read_csv(csv_filename)\n",
						"\n",
						"    fig.add_trace(go.Scatter(\n",
						"        x=df_repo['timestamp'],\n",
						"        y=df_repo['open_issues_count'],\n",
						"        mode='lines+markers',\n",
						"        name='Open Issues',\n",
						"        hovertemplate='Open Issues: %{y}<br>Time: %{x}<extra></extra>'\n",
						"    ))\n"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"for REPO in REPOS:\n",
						"    plot_repo_metrics(f\"{OWNER}/{REPO}\")"
					]
				}
			]
				.map(fromJupyterCell)
			,
			[
				{
					"cell_type": "markdown",
					"metadata": {},
					"source": [
						"# Fetch Data"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import requests\n",
						"import sys\n",
						"import pandas as pd\n",
						"import matplotlib.pyplot as plt\n",
						"from cachetools import cached, TTLCache\n",
						"from dotenv import load_dotenv"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import os\n",
						"\n",
						"GITHUB_API_URL = \"https://api.github.com\"\n",
						"REPO_OWNER = \"microsoft\"\n",
						"ISSUE = 123\n",
						"load_dotenv()\n",
						"TOKEN = os.getenv(\"GH\")\n",
						"\n",
						"headers = {\n",
						"    \"Authorization\": f\"Bearer {TOKEN}\",\n",
						"    \"Accept\": \"application/vnd.github+json\"\n",
						"}"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import csv\n",
						"from datetime import datetime\n",
						"import os\n",
						"\n",
						"def append_information(repo, open_issues_count, ):\n",
						"    file_exists = os.path.isfile(filename)\n",
						"    timestamp = datetime.utcnow().isoformat()\n",
						"\n",
						"    # Open the file in append mode\n",
						"    with open(filename, 'a', newline='', encoding='utf-8') as csvfile:\n",
						"        writer = csv.writer(csvfile)\n",
						"\n",
						"        # If the file didn't exist before, write the header row\n",
						"        if not file_exists:\n",
						"            writer.writerow([\n",
						"                \"timestamp\",\n",
						"                \"repo\",\n",
						"                \"open_issues_count\",\n",
						"                \"open_prs_count\",\n",
						"                \"bug_issues_count\",\n",
						"                \"feature_request_issues_count\",\n",
						"                \"invalid_issues_count\"\n",
						"            ])\n"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"def get_repo_issue_stats(owner, repo, token):\n",
						"    # Common GitHub GraphQL endpoint\n",
						"    url = \"https://api.github.com/graphql\"\n",
						"    headers = {\n",
						"        \"Authorization\": f\"Bearer {token}\",\n",
						"        \"Accept\": \"application/vnd.github.v4+json\"\n",
						"    }\n",
						"\n",
						"    # First query: Get various issue/PR counts from the repository object\n",
						"    query_repo = \"\"\"\n",
						"    query ($owner: String!, $name: String!) {\n",
						"      repository(owner: $owner, name: $name) {\n",
						"        issues(states: OPEN) {\n",
						"          totalCount\n",
						"        }\n",
						"        pullRequests(states: OPEN) {\n",
						"          totalCount\n",
						"        }\n",
						"        bugIssues: issues(states: OPEN, labels: [\"bug\"]) {\n",
						"          totalCount\n",
						"        }\n",
						"        featureRequestIssues: issues(states: OPEN, labels: [\"feature-request\"]) {\n",
						"          totalCount\n",
						"        }\n",
						"      }\n",
						"    }\n",
						"    \"\"\"\n"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"for REPO in REPOS:\n",
						"    stats = get_repo_issue_stats(OWNER, REPO, TOKEN)\n",
						"    print(stats)\n",
						"\n",
						"    append_metrics_to_csv(\n",
						"        filename=\"repo_metrics.csv\",\n",
						"        repo=f\"{OWNER}/{REPO}\",\n",
						"        open_issues_count=stats['open_issues_count'],\n",
						"        open_prs_count=stats['open_prs_count'],\n",
						"        bug_issues_count=stats['bug_issues_count'],\n",
						"        feature_request_issues_count=stats['feature_request_issues_count'],\n",
						"        invalid_issues_count=stats['invalid_issues_count']\n",
						"    )"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {},
					"source": [
						"# Plot"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"import pandas as pd\n",
						"import plotly.graph_objs as go\n",
						"from plotly.subplots import make_subplots\n",
						"\n",
						"def plot_repo_metrics(repo, csv_filename=\"repo_metrics.csv\"):\n",
						"    # Load the CSV data\n",
						"    df = pd.read_csv(csv_filename)\n",
						"\n",
						"    fig.add_trace(go.Scatter(\n",
						"        x=df_repo['timestamp'],\n",
						"        y=df_repo['open_issues_count'],\n",
						"        mode='lines+markers',\n",
						"        name='Open Issues',\n",
						"        hovertemplate='Open Issues: %{y}<br>Time: %{x}<extra></extra>'\n",
						"    ))\n"
					]
				},
				{
					"cell_type": "markdown",
					"metadata": {},
					"source": [
						"## Header"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"for REPO in REPOS:\n",
						"    plot_repo_metrics(f\"{OWNER}/{REPO}\")"
					]
				},
				{
					"cell_type": "code",
					"execution_count": null,
					"metadata": {},
					"outputs": [],
					"source": [
						"plot_repo_metrics(f\"{OWNER}/vscode\")"
					]
				}
			].map(fromJupyterCell)
		);

		assert.deepStrictEqual(mapping, [
			{ modified: 0, original: -1 },
			{ modified: 1, original: 0 },
			{ modified: 2, original: 1 },
			{ modified: 3, original: 2 },
			{ modified: 4, original: 3 },
			{ modified: 5, original: 4 },
			{ modified: 6, original: 5 },
			{ modified: 7, original: 6 },
			{ modified: 8, original: -1 },
			{ modified: 9, original: 7 },
			{ modified: 10, original: -1 }
		]);
		assert.deepStrictEqual(diff.cellsDiff.changes, [
			{ originalStart: 0, originalLength: 0, modifiedStart: 0, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 0, originalLength: 1, modifiedStart: 1, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 1, originalLength: 1, modifiedStart: 2, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 7, originalLength: 0, modifiedStart: 8, modifiedLength: 1 } satisfies IDiffChange,
			{ originalStart: 8, originalLength: 0, modifiedStart: 10, modifiedLength: 1 } satisfies IDiffChange,
		]);


	});

	let uriCounter = 0;
	async function mapCells(original: MockNotebookCell[], modified: MockNotebookCell[]) {
		const mapping = matchCellBasedOnSimilarties(modified.map(cellToDto), original.map(cellToDto)).map(mapping => ({ modified: mapping.modified, original: mapping.original }));
		const originalUri = URI.parse(`file://original${uriCounter++}.ipynb`);
		const modifiedUri = URI.parse(`file://modified${uriCounter++}.ipynb`);
		worker.$acceptNewModel(originalUri.toString(), {}, {}, original.map((cell, index) => cellToMainCellTdo(originalUri, cell, index)));
		worker.$acceptNewModel(modifiedUri.toString(), {}, {}, modified.map((cell, index) => cellToMainCellTdo(modifiedUri, cell, index)));
		const diff = await worker.$computeDiff(originalUri.toString(), modifiedUri.toString());
		diff.cellsDiff.changes = diff.cellsDiff.changes.map(change => {
			return {
				originalStart: change.originalStart,
				originalLength: change.originalLength,
				modifiedStart: change.modifiedStart,
				modifiedLength: change.modifiedLength,
			};
		});
		return { mapping, diff };
	}


	function cellToDto(cell: MockNotebookCell): { getValue(): string; getLinesContent(): string[]; cellKind: CellKind } {
		return {
			cellKind: cell[2],
			getValue: () => cell[0].join('\n'),
			getLinesContent: () => cell[0]
		};
	}

	function cellToMainCellTdo(notebookUri: URI, cell: MockNotebookCell, index: number): IMainCellDto {
		return {
			source: cell[0],
			language: cell[1],
			cellKind: cell[2],
			outputs: cell[3] || [],
			metadata: cell[4],
			eol: '\n',
			handle: index,
			url: URI.from({ scheme: 'file', path: notebookUri.path, fragment: `cell/${index}` }).toString(),
			versionId: 0,
			internalMetadata: undefined
		} satisfies IMainCellDto;

	}
	type MockNotebookCell = [
		lines: string[],
		lang: string,
		kind: CellKind,
		output?: IOutputDto[],
		metadata?: NotebookCellMetadata,
	];

	function fromJupyterCell(cell: { cell_type: string; source: string[] }): MockNotebookCell {
		return [
			cell.source,
			cell.cell_type === 'code' ? 'python' : 'markdown',
			cell.cell_type === 'code' ? CellKind.Code : CellKind.Markup,
		];
	}
});


