feat: added similar face search using embeddings

Signed-off-by: Varun Raj <mailme@varunraj.in>
This commit is contained in:
Varun Raj
2024-08-23 23:18:40 +05:30
parent cc04bf2e34
commit d0a6d1d9b7
10 changed files with 513 additions and 63 deletions
+4 -4
View File
@@ -1,7 +1,7 @@
IMMICH_URL="" # Immich URL
IMMICH_API_KEY="" # Immich API Key
DB_USERNAME="" # Postgress Database Username
DB_PASSWORD = "" # Postgres Database Password
DB_HOST = "" # Postgres Host (IP address or hostname of the database)
DB_PORT = "" # Postgres Port number (Default: 5432)
DB_DATABASE_NAME = "" # Name of the database
DB_PASSWORD="" # Postgres Database Password
DB_HOST="" # Postgres Host (IP address or hostname of the database)
DB_PORT="" # Postgres Port number (Default: 5432)
DB_DATABASE_NAME="" # Name of the database
+265 -44
View File
@@ -1,44 +1,75 @@
"use client";
import { Button } from "@/components/ui/button";
import {
Command,
CommandEmpty,
CommandGroup,
CommandInput,
CommandItem,
CommandList,
} from "@/components/ui/command";
import { Dialog, DialogContent, DialogTrigger } from "@/components/ui/dialog";
import { Merge } from "lucide-react";
import { mergePerson, searchPeople } from "@/handlers/api/people.handler";
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
DialogTrigger,
} from "@/components/ui/dialog";
import { SearchSlash, X } from "lucide-react";
import {
listSimilarFaces,
mergePerson,
searchPeople,
} from "@/handlers/api/people.handler";
import { IPerson } from "@/types/person";
import { set } from "date-fns";
import { CommandLoading } from "cmdk";
import { Avatar } from "../ui/avatar";
import { useEffect, useRef, useState } from "react";
import { useEffect, useMemo, useRef, useState } from "react";
import { useToast } from "../ui/use-toast";
import { Input } from "../ui/input";
import Loader from "../ui/loader";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuTrigger,
} from "../ui/dropdown-menu";
import { CaretDownIcon, FaceIcon } from "@radix-ui/react-icons";
import FaceThumbnail from "./merge/FaceThumbnail";
import ErrorBlock from "../shared/ErrorBlock";
interface PersonMergeDropdownProps {
person: IPerson;
onRemove?: (person: IPerson) => void;
}
export function PersonMergeDropdown({
person,
onRemove,
}: PersonMergeDropdownProps) {
const [searchedPeople, setSearchedPeople] = useState<IPerson[]>([]);
const [searchedPeople, setSearchedPeople] = useState<IPerson[] | null>(null);
const [similarPeople, setSimilarPeople] = useState<IPerson[]>([]);
const [open, setOpen] = useState(false);
const [loading, setLoading] = useState(false);
const [similarLoading, setSimilarLoading] = useState(true);
const searchTimer = useRef<NodeJS.Timeout>();
const [selectedPeople, setSelectedPeople] = useState<IPerson[]>([]);
const [primaryPerson, setPrimaryPerson] = useState<IPerson>(person);
const [merging, setMerging] = useState(false);
const { toast } = useToast();
const selectedIds = useMemo(
() => selectedPeople.map((p) => p.id),
[selectedPeople]
);
const handleSearch = (value: string) => {
if (searchTimer.current) {
clearTimeout(searchTimer.current);
}
searchTimer.current = setTimeout(() => {
if (!value || !value.trim().length) {
setSearchedPeople(null);
return;
}
setLoading(true);
if (!value) {
setSearchedPeople([]);
@@ -54,23 +85,144 @@ export function PersonMergeDropdown({
}, 500);
};
const fetchSuggestions = () => {
return listSimilarFaces(person.id)
.then(setSimilarPeople)
.catch(() => {
toast({
title: "Error",
description: "Failed to fetch similar people",
});
})
.finally(() => {
setSimilarLoading(false);
});
};
const handleSelect = (value: IPerson) => {
setOpen(false);
return mergePerson(person.id, value.id)
const isAlreadySelected = selectedPeople.some((p) => p.id === value.id);
if (isAlreadySelected) {
setSelectedPeople(selectedPeople.filter((p) => p.id !== value.id));
}
if (selectedPeople.length >= 5) {
toast({
title: "Error",
description: "You can only merge 5 people at a time",
});
return;
} else {
setSelectedPeople([...selectedPeople, value]);
}
if (primaryPerson.name.length === 0 && value.name.length > 0) {
setPrimaryPerson(value);
}
// setOpen(false);
// return mergePerson(person.id, value.id)
// .then(() => {
// onRemove?.(person);
// })
// .then(() => {
// toast({
// title: "Success",
// description: "Person merged successfully",
// });
// });
};
const handleMerge = () => {
if (selectedPeople.length === 0) {
return;
}
const personIds = selectedPeople.map((p) => p.id);
setMerging(true);
return mergePerson(person.id, personIds)
.then(() => {
onRemove?.(person);
})
.then(() => {
setOpen(false);
toast({
title: "Success",
description: "Person merged successfully",
description: "People merged successfully",
});
})
.catch(() => {
toast({
title: "Error",
description: "Failed to merge people",
});
})
.finally(() => {
setMerging(false);
});
};
const handleRemove = (value: IPerson) => {
const newSelected = selectedPeople.filter((p) => p.id !== value.id);
if (primaryPerson.id === value.id) {
const newPrimary = newSelected[0];
setPrimaryPerson(newPrimary || person);
}
setSelectedPeople(newSelected);
};
useEffect(() => {
if (open) handleSearch("");
}, [open]);
if (open && !similarPeople.length) fetchSuggestions();
}, [open, person.id]);
const renderPeopleList = (people: IPerson[], title: string) => {
return (
<div className="max-h-[400px] min-h-[400px] overflow-y-auto">
<p className="pb-2 uppercase">{title}</p>
<div className="grid grid-cols-4 gap-2">
{people.map((person) => (
<FaceThumbnail
key={person.id}
person={person}
onSelect={handleSelect}
selected={selectedIds.includes(person.id)}
/>
))}
</div>
</div>
);
};
const renderContent = () => {
if (loading || similarLoading) {
return (
<div className="min-h-[400px] flex flex-col items-center justify-center">
<Loader />
</div>
);
}
if (searchedPeople && searchedPeople?.length === 0) {
return (
<div className="min-h-[400px] flex flex-col items-center justify-center">
<p>No Results Found</p>
</div>
);
}
if (searchedPeople && searchedPeople?.length > 0) {
return renderPeopleList(searchedPeople, "Search Results");
}
if (similarPeople && similarPeople.length === 0) {
return (
<ErrorBlock
title="No Similar People Found"
description="Please search for the person"
icon={<SearchSlash />}
/>
);
}
if (similarPeople && similarPeople.length > 0) {
return renderPeopleList(similarPeople, "Similar People");
}
return <p className="text-xs">Please search for people</p>;
};
return (
<div className="flex items-center space-x-4">
@@ -80,35 +232,104 @@ export function PersonMergeDropdown({
Merge
</Button>
</DialogTrigger>
<DialogContent className="p-0">
<Command>
<CommandInput
placeholder="Change status..."
onValueChange={handleSearch}
/>
<CommandList>
{loading && <CommandLoading />}
<CommandGroup heading="People">
{searchedPeople.map((person) => (
<CommandItem
key={person.id}
value={person.name}
onSelect={() => {
handleSelect(person);
}}
className="flex items-center gap-2"
>
<DialogContent>
<DialogHeader>
<DialogTitle>
Merge {person.name ? person.name : "Untagged Person"}
</DialogTitle>
<DialogDescription>
<div className="flex items-center gap-2">
<p>
{selectedPeople.length > 0
? `Merging ${selectedPeople.length} people to ${person.name}`
: "Search and select people to merge with"}
</p>
<DropdownMenu>
<DropdownMenuTrigger className="flex gap-1 items-center">
<Avatar
src={person.thumbnailPath}
alt={person.name}
className="w-6 h-6"
src={primaryPerson.thumbnailPath}
alt={primaryPerson.name}
className="w-4 h-4"
/>
<span>{person.name}</span>
</CommandItem>
))}
</CommandGroup>
</CommandList>
</Command>
<p>{primaryPerson?.name || "Untagged Person"}</p>
<CaretDownIcon />
</DropdownMenuTrigger>
<DropdownMenuContent>
{[person, ...selectedPeople].map((person) => (
<DropdownMenuItem
key={person.id}
onSelect={() => setPrimaryPerson(person)}
className="flex gap-1"
>
<Avatar
src={person.thumbnailPath}
alt={person.name}
className="w-6 h-6"
/>
<span>{person.name || "Untagged Person"}</span>
</DropdownMenuItem>
))}
</DropdownMenuContent>
</DropdownMenu>
</div>
</DialogDescription>
</DialogHeader>
<Input
placeholder="Search people..."
onChange={(e) => {
handleSearch(e.target.value);
}}
/>
{selectedPeople.length > 0 ? (
<div className="flex flex-nowrap overflow-x-auto gap-2 py-2">
{selectedPeople.map((person) => (
<div
key={person.id}
className="flex border px-1 items-center gap-1 bg-zinc-900 rounded-lg p-1"
>
<Avatar
src={person.thumbnailPath}
alt={person.name}
className="w-6 h-6"
/>
<p className="text-xs text-nowrap">
{person.name ? person.name : <span>Untagged person</span>}
</p>
<button
className="rounded-full hover:bg-zinc-800 p-1"
onClick={() => handleRemove(person)}
>
<X className="w-4 h-4" />
</button>
</div>
))}
</div>
) : (
<div>
<p className="py-4 text-sm text-zinc-500">No Selections Yet</p>
</div>
)}
{renderContent()}
<DialogFooter>
<Button
onClick={() => {
setOpen(false);
setSelectedPeople([]);
setPrimaryPerson(person);
}}
variant="outline"
>
Cancel
</Button>
<Button
onClick={handleMerge}
disabled={selectedPeople.length === 0 || merging}
>
Merge {selectedPeople.length + 1} People
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</div>
@@ -0,0 +1,35 @@
import { Avatar } from "@/components/ui/avatar";
import { cn } from "@/lib/utils";
import { IPerson } from "@/types/person";
interface FaceThumbnailProps {
person: IPerson;
onSelect: (person: IPerson) => void;
selected?: boolean;
}
const FaceThumbnail = ({ person, onSelect, selected }: FaceThumbnailProps) => {
return (
<div
className={cn(
"col-span-1 flex flex-col items-center p-2 border-2 text-center rounded-lg hover:bg-zinc-900 gap-2 justify-center",
{
"border-zinc-500": selected,
}
)}
onClick={() => onSelect(person)}
>
<Avatar
src={person.thumbnailPath}
alt={person.name}
className="w-20 h-20"
/>
{person.name ? (
<p className="text-xs">{person.name}</p>
) : (
<p className="text-xs italic text-gray-400">Untagged Person</p>
)}
</div>
);
};
export default FaceThumbnail;
+18
View File
@@ -0,0 +1,18 @@
import React from 'react'
interface ErrorBlockProps {
icon?: React.ReactNode;
title?: string;
description?: string;
}
export default function ErrorBlock(
{ icon, title, description }: ErrorBlockProps
) {
return (
<div className="flex flex-col items-center justify-center w-full h-full space-y-4 py-10">
{icon}
{title && <h2 className="text-xl font-semibold">{title}</h2>}
{description && <p className="text-gray-500">{description}</p>}
</div>
)
}
+162
View File
@@ -0,0 +1,162 @@
import * as React from "react"
import {
CaretSortIcon,
CheckIcon,
ChevronDownIcon,
ChevronUpIcon,
} from "@radix-ui/react-icons"
import * as SelectPrimitive from "@radix-ui/react-select"
import { cn } from "@/lib/utils"
const Select = SelectPrimitive.Root
const SelectGroup = SelectPrimitive.Group
const SelectValue = SelectPrimitive.Value
const SelectTrigger = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.Trigger>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.Trigger>
>(({ className, children, ...props }, ref) => (
<SelectPrimitive.Trigger
ref={ref}
className={cn(
"flex h-9 w-full items-center justify-between whitespace-nowrap rounded-md border border-input bg-transparent px-3 py-2 text-sm shadow-sm ring-offset-background placeholder:text-muted-foreground focus:outline-none focus:ring-1 focus:ring-ring disabled:cursor-not-allowed disabled:opacity-50 [&>span]:line-clamp-1",
className
)}
{...props}
>
{children}
<SelectPrimitive.Icon asChild>
<CaretSortIcon className="h-4 w-4 opacity-50" />
</SelectPrimitive.Icon>
</SelectPrimitive.Trigger>
))
SelectTrigger.displayName = SelectPrimitive.Trigger.displayName
const SelectScrollUpButton = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.ScrollUpButton>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.ScrollUpButton>
>(({ className, ...props }, ref) => (
<SelectPrimitive.ScrollUpButton
ref={ref}
className={cn(
"flex cursor-default items-center justify-center py-1",
className
)}
{...props}
>
<ChevronUpIcon />
</SelectPrimitive.ScrollUpButton>
))
SelectScrollUpButton.displayName = SelectPrimitive.ScrollUpButton.displayName
const SelectScrollDownButton = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.ScrollDownButton>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.ScrollDownButton>
>(({ className, ...props }, ref) => (
<SelectPrimitive.ScrollDownButton
ref={ref}
className={cn(
"flex cursor-default items-center justify-center py-1",
className
)}
{...props}
>
<ChevronDownIcon />
</SelectPrimitive.ScrollDownButton>
))
SelectScrollDownButton.displayName =
SelectPrimitive.ScrollDownButton.displayName
const SelectContent = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.Content>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.Content>
>(({ className, children, position = "popper", ...props }, ref) => (
<SelectPrimitive.Portal>
<SelectPrimitive.Content
ref={ref}
className={cn(
"relative z-50 max-h-96 min-w-[8rem] overflow-hidden rounded-md border bg-popover text-popover-foreground shadow-md data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2",
position === "popper" &&
"data-[side=bottom]:translate-y-1 data-[side=left]:-translate-x-1 data-[side=right]:translate-x-1 data-[side=top]:-translate-y-1",
className
)}
position={position}
{...props}
>
<SelectScrollUpButton />
<SelectPrimitive.Viewport
className={cn(
"p-1",
position === "popper" &&
"h-[var(--radix-select-trigger-height)] w-full min-w-[var(--radix-select-trigger-width)]"
)}
>
{children}
</SelectPrimitive.Viewport>
<SelectScrollDownButton />
</SelectPrimitive.Content>
</SelectPrimitive.Portal>
))
SelectContent.displayName = SelectPrimitive.Content.displayName
const SelectLabel = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.Label>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.Label>
>(({ className, ...props }, ref) => (
<SelectPrimitive.Label
ref={ref}
className={cn("px-2 py-1.5 text-sm font-semibold", className)}
{...props}
/>
))
SelectLabel.displayName = SelectPrimitive.Label.displayName
const SelectItem = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.Item>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.Item>
>(({ className, children, ...props }, ref) => (
<SelectPrimitive.Item
ref={ref}
className={cn(
"relative flex w-full cursor-default select-none items-center rounded-sm py-1.5 pl-2 pr-8 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50",
className
)}
{...props}
>
<span className="absolute right-2 flex h-3.5 w-3.5 items-center justify-center">
<SelectPrimitive.ItemIndicator>
<CheckIcon className="h-4 w-4" />
</SelectPrimitive.ItemIndicator>
</span>
<SelectPrimitive.ItemText>{children}</SelectPrimitive.ItemText>
</SelectPrimitive.Item>
))
SelectItem.displayName = SelectPrimitive.Item.displayName
const SelectSeparator = React.forwardRef<
React.ElementRef<typeof SelectPrimitive.Separator>,
React.ComponentPropsWithoutRef<typeof SelectPrimitive.Separator>
>(({ className, ...props }, ref) => (
<SelectPrimitive.Separator
ref={ref}
className={cn("-mx-1 my-1 h-px bg-muted", className)}
{...props}
/>
))
SelectSeparator.displayName = SelectPrimitive.Separator.displayName
export {
Select,
SelectGroup,
SelectValue,
SelectTrigger,
SelectContent,
SelectLabel,
SelectItem,
SelectSeparator,
SelectScrollUpButton,
SelectScrollDownButton,
}
+1
View File
@@ -5,6 +5,7 @@ export const GET_ME_PATH = BASE_PROXY_ENDPOINT + "/users/me";
export const LIST_PEOPLE_PATH = BASE_API_ENDPOINT + "/people/list";
export const SEARCH_PEOPLE_PATH = BASE_PROXY_ENDPOINT + "/search/person";
export const SIMILAR_FACES_PATH = (id: string) => BASE_API_ENDPOINT + "/people/" + id + "/similar-faces";
export const PERSON_THUBNAIL_PATH = (id: string) => BASE_PROXY_ENDPOINT + "/thumbnail/" + id;
export const UPDATE_PERSON_PATH = (id: string) => BASE_PROXY_ENDPOINT + "/people/" + id;
export const MERGE_PERSON_PATH = (id: string) => BASE_PROXY_ENDPOINT + "/people/" + id + "/merge";
+7 -3
View File
@@ -1,4 +1,4 @@
import { LIST_PEOPLE_PATH, MERGE_PERSON_PATH, SEARCH_PEOPLE_PATH, UPDATE_PERSON_PATH } from "@/config/routes"
import { LIST_PEOPLE_PATH, MERGE_PERSON_PATH, SEARCH_PEOPLE_PATH, SIMILAR_FACES_PATH, UPDATE_PERSON_PATH } from "@/config/routes"
import { cleanUpPerson } from "@/helpers/person.helper";
import API from "@/lib/api"
import { IPeopleListResponse, IPerson } from "@/types/person"
@@ -36,6 +36,10 @@ export const searchPeople = (name: string) => {
return API.get(SEARCH_PEOPLE_PATH, { name }).then((response) => response.map(cleanUpPerson));
}
export const mergePerson = (id: string, targetId: string) => {
return API.post(MERGE_PERSON_PATH(id), { ids: [ targetId ] })
export const mergePerson = (id: string, targetIds: string[]) => {
return API.post(MERGE_PERSON_PATH(id), { ids: targetIds })
}
export const listSimilarFaces = (id: string) => {
return API.get(SIMILAR_FACES_PATH(id)).then((response) => response.map(cleanUpPerson));
}
@@ -53,6 +53,7 @@ export default async function handler(
isNotNull(assets.createdAt),
gte(exif.dateTimeOriginal, startDateDate),
lte(exif.dateTimeOriginal, endDateDate),
eq(assets.type, "VIDEO"),
eq(assets.ownerId, currentUser.id),
));
@@ -22,6 +22,7 @@ export default async function handler(
.leftJoin(exif, eq(exif.assetId, assets.id))
.where(and(
isNull(exif.latitude),
eq(assets.type, "VIDEO"),
eq(assets.ownerId, currentUser.id),
))
.groupBy(sql`DATE(${exif.dateTimeOriginal})`)
+19 -12
View File
@@ -1,10 +1,11 @@
// Next.js API route support: https://nextjs.org/docs/api-routes/introduction
import { CHART_COLORS } from "@/config/constants/chart.constant";
import { db } from "@/config/db";
import { getCurrentUser } from "@/handlers/serverUtils/user.utils";
import { stringToBoolean } from "@/helpers/data.helper";
import { assetFaces, assets, exif, person } from "@/schema";
import { faceSearch } from "@/schema/faceSearch.schema";
import { cosineDistance, desc, eq, gt, ne, sql } from "drizzle-orm";
import { and, cosineDistance, desc, eq, gt, ne, sql } from "drizzle-orm";
import type { NextApiRequest, NextApiResponse } from "next";
type ISortField = "assetCount" | "updatedAt" | "createdAt";
@@ -33,11 +34,13 @@ export default async function handler(
sortOrder = "desc",
} = req.query as any as IQuery;
const currentUser = await getCurrentUser();
const personRecords = await db
.select()
.from(person)
.where(eq(person.id, id))
.limit(1);
const personRecord = personRecords?.[0];
if (!personRecord) {
return res.status(404).json({
@@ -57,11 +60,12 @@ export default async function handler(
error: "Person has no face",
});
}
console.log(assetFaceRecord.id);
const faceSearchRecords = await db
.select()
.from(faceSearch)
.where(ne(faceSearch.faceId, assetFaceRecord.id))
.where(eq(faceSearch.faceId, assetFaceRecord.id))
.limit(1);
const faceSearchRecord = faceSearchRecords?.[0];
@@ -77,28 +81,31 @@ export default async function handler(
)})`;
const people = await db
.select({
.selectDistinctOn([person.id], {
id: person.id,
name: person.name,
thumbnailPath: person.thumbnailPath,
birthDate: person.birthDate,
isHidden: person.isHidden,
updatedAt: person.updatedAt,
assetId: assetFaces.assetId,
assetId: assetFaces.id,
faceSearch: faceSearch.faceId,
similarity,
})
.from(faceSearch)
.leftJoin(assetFaces, eq(assetFaces.id, faceSearch.faceId))
.leftJoin(person, eq(person.id, assetFaces.personId))
.where(gt(similarity, 0.5))
.limit(4);
.innerJoin(person, eq(person.id, assetFaces.personId))
console.log(people);
.where(
and(
ne(person.id, id),
eq(person.ownerId, currentUser.id),
gt(similarity, 0.5)
)
)
.limit(12);
return res.status(200).json({
people,
total: 10,
});
return res.status(200).json(people);
} catch (error: any) {
res.status(500).json({
error: error?.message,