Add support for video streaming, video thumbnail generation, video search functionality, and video tagging

pull/550/head
zanllp 2024-03-22 06:29:48 +08:00
parent b72e251cee
commit 0f739ae155
11 changed files with 213 additions and 54 deletions

View File

@ -2,4 +2,6 @@ fastapi
uvicorn
piexif
python-dotenv
Pillow
Pillow
imageio
av

View File

@ -1,14 +1,16 @@
from datetime import datetime, timedelta
import os
from pathlib import Path
import shutil
import sqlite3
from scripts.iib.tool import (
comfyui_exif_data_to_str,
get_comfyui_exif_data,
human_readable_size,
is_img_created_by_comfyui,
is_img_created_by_comfyui_with_webui_gen_info,
is_valid_image_path,
is_valid_media_path,
temp_path,
read_sd_webui_gen_info_from_image,
get_formatted_date,
@ -28,12 +30,12 @@ from scripts.iib.tool import (
is_secret_key_required,
open_file_with_default_app
)
from fastapi import FastAPI, HTTPException, Response
from fastapi import FastAPI, HTTPException, Header, Response
from fastapi.staticfiles import StaticFiles
import asyncio
from typing import List, Optional
from pydantic import BaseModel
from fastapi.responses import FileResponse, JSONResponse
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from PIL import Image
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
@ -53,6 +55,7 @@ from scripts.iib.db.update_image_data import update_image_data, rebuild_image_in
from scripts.iib.logger import logger
from scripts.iib.seq import seq
import urllib.parse
from scripts.iib.fastapi_video import range_requests_response
index_html_path = os.path.join(cwd, "vue/dist/index.html") # 在app.py也被使用
@ -493,7 +496,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
encoded_filename = urllib.parse.quote(disposition.encode('utf-8'))
headers['Content-Disposition'] = f"attachment; filename*=UTF-8''{encoded_filename}"
if is_path_under_parents(filename) and is_valid_image_path(
if is_path_under_parents(filename) and is_valid_media_path(
filename
): # 认为永远不变,不要协商缓存了试试
headers[
@ -508,6 +511,54 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
media_type=media_type,
headers=headers,
)
@app.get(api_base + "/stream_video", dependencies=[Depends(verify_secret)])
async def stream_video(path: str, request: Request):
import mimetypes
print(path)
print(request)
media_type, _ = mimetypes.guess_type(path)
return range_requests_response(
request, file_path=path, content_type=media_type
)
@app.get(api_base + "/video_cover", dependencies=[Depends(verify_secret)])
async def video_cover(path: str, t: str):
check_path_trust(path)
if not temp_path:
return
# 生成缓存文件的路径
hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest()
hash = hash_dir
cache_dir = os.path.join(temp_path, "iib_cache", "video_cover", hash_dir)
cache_path = os.path.join(cache_dir, "cover.webp")
# 如果缓存文件存在,则直接返回该文件
if os.path.exists(cache_path):
print("命中缓存")
return FileResponse(
cache_path,
media_type="image/webp",
headers={"Cache-Control": "max-age=31536000", "ETag": hash},
)
# 如果缓存文件不存在,则生成缩略图并保存
import imageio.v3 as iio
frame = iio.imread(
path,
index=16,
plugin="pyav",
)
os.makedirs(cache_dir, exist_ok=True)
iio.imwrite(cache_path,frame, extension=".webp")
# 返回缓存文件
return FileResponse(
cache_path,
media_type="image/webp",
headers={"Cache-Control": "max-age=31536000", "ETag": hash},
)
@app.post(api_base + "/send_img_path", dependencies=[Depends(verify_secret)])
async def api_set_send_img_path(path: str):
@ -728,7 +779,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
@app.get(db_api_base + "/img_selected_custom_tag", dependencies=[Depends(verify_secret)])
async def get_img_selected_custom_tag(path: str):
path = os.path.normpath(path)
if not is_valid_image_path(path):
if not is_valid_media_path(path):
return []
conn = DataBase.get_conn()
update_extra_paths(conn)

View File

@ -6,8 +6,9 @@ from PIL import Image
from scripts.iib.tool import (
read_sd_webui_gen_info_from_image,
parse_generation_parameters,
is_valid_image_path,
is_valid_media_path,
get_modified_date,
get_video_type,
is_dev,
get_comfyui_exif_data,
comfyui_exif_data_to_str,
@ -22,6 +23,8 @@ from scripts.iib.logger import logger
def get_exif_data(file_path):
info = ''
params = None
if get_video_type(file_path):
return params, info
try:
with Image.open(file_path) as img:
if is_img_created_by_comfyui(img):
@ -62,7 +65,7 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
if os.path.isdir(file_path):
process_folder(file_path)
elif is_valid_image_path(file_path):
elif is_valid_media_path(file_path):
img = DbImg.get(conn, file_path)
if is_rebuild:
parsed_params, info = get_exif_data(file_path)

View File

@ -0,0 +1,74 @@
import os
from typing import BinaryIO
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import StreamingResponse
def send_bytes_range_requests(
file_obj: BinaryIO, start: int, end: int, chunk_size: int = 10_000
):
"""Send a file in chunks using Range Requests specification RFC7233
`start` and `end` parameters are inclusive due to specification
"""
with file_obj as f:
f.seek(start)
while (pos := f.tell()) <= end:
read_size = min(chunk_size, end + 1 - pos)
yield f.read(read_size)
def _get_range_header(range_header: str, file_size: int) -> tuple[int, int]:
def _invalid_range():
return HTTPException(
status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
detail=f"Invalid request range (Range:{range_header!r})",
)
try:
h = range_header.replace("bytes=", "").split("-")
start = int(h[0]) if h[0] != "" else 0
end = int(h[1]) if h[1] != "" else file_size - 1
except ValueError:
raise _invalid_range()
if start > end or start < 0 or end > file_size - 1:
raise _invalid_range()
return start, end
def range_requests_response(
request: Request, file_path: str, content_type: str
):
"""Returns StreamingResponse using Range Requests of a given file"""
file_size = os.stat(file_path).st_size
range_header = request.headers.get("range")
headers = {
"content-type": content_type,
"accept-ranges": "bytes",
"content-encoding": "identity",
"content-length": str(file_size),
"access-control-expose-headers": (
"content-type, accept-ranges, content-length, "
"content-range, content-encoding"
),
}
start = 0
end = file_size - 1
status_code = status.HTTP_200_OK
if range_header is not None:
start, end = _get_range_header(range_header, file_size)
size = end - start + 1
headers["content-length"] = str(size)
headers["content-range"] = f"bytes {start}-{end}/{file_size}"
status_code = status.HTTP_206_PARTIAL_CONTENT
return StreamingResponse(
send_bytes_range_requests(open(file_path, mode="rb"), start, end),
headers=headers,
status_code=status_code,
)

View File

@ -167,9 +167,17 @@ def convert_to_bytes(file_size_str):
return int(size)
else:
raise ValueError(f"Invalid file size string '{file_size_str}'")
def get_video_type(file_path):
video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv']
file_extension = file_path[file_path.rfind('.'):].lower()
if file_extension in video_extensions:
return file_extension[1:]
else:
return None
def is_valid_image_path(path):
def is_valid_media_path(path):
"""
判断给定的路径是否是图像文件
"""
@ -178,7 +186,7 @@ def is_valid_image_path(path):
return False
if not os.path.isfile(abs_path): # 判断是否是文件
return False
if not imghdr.what(abs_path): # 判断是否是图像文件
if not imghdr.what(abs_path) and not get_video_type(abs_path): # 判断是否是图像文件
return False
return True

View File

@ -2,7 +2,7 @@
import type { Tag } from '@/api/db'
import type { FileNodeInfo } from '@/api/files'
import type { MenuInfo } from 'ant-design-vue/lib/menu/src/interface'
import { isImageFile } from '@/util'
import { isMediaFile } from '@/util'
import { StarFilled, StarOutlined } from '@/icon'
import { useGlobalStore } from '@/store/useGlobalStore'
import { computed } from 'vue'
@ -35,7 +35,7 @@ const tags = computed(() => {
<a-menu-item key="openWithWalkMode">{{ $t('openWithWalkMode') }}</a-menu-item>
</template>
<template v-if="file.type === 'file'">
<template v-if="isImageFile(file.name)">
<template v-if="isMediaFile(file.name)">
<a-menu-item key="viewGenInfo">{{ $t('viewGenerationInfo') }}</a-menu-item>
<a-menu-divider />
<template v-if="global.conf?.launch_mode !== 'server'">

View File

@ -4,16 +4,17 @@ import { useGlobalStore } from '@/store/useGlobalStore'
import { fallbackImage, ok } from 'vue3-ts-util'
import type { FileNodeInfo } from '@/api/files'
import { isImageFile, isVideoFile } from '@/util'
import { toImageThumbnailUrl, toRawFileUrl } from '@/util/file'
import { toImageThumbnailUrl, toVideoCoverUrl, toRawFileUrl } from '@/util/file'
import type { MenuInfo } from 'ant-design-vue/lib/menu/src/interface'
import { computed } from 'vue'
import ContextMenu from './ContextMenu.vue'
import ChangeIndicator from './ChangeIndicator.vue'
import { useTagStore } from '@/store/useTagStore'
import { CloseCircleOutlined, StarFilled, StarOutlined, PlayCircleFilled } from '@/icon'
import { CloseCircleOutlined, StarFilled, StarOutlined } from '@/icon'
import { Tag } from '@/api/db'
import { openVideoModal } from './functionalCallableComp'
import type { GenDiffInfo } from '@/api/files'
import { play } from '@/icon'
const global = useGlobalStore()
const tagStore = useTagStore()
@ -35,19 +36,23 @@ const props = withDefaults(
enableChangeIndicator?: boolean
extraTags?: Tag[]
}>(),
{ selected: false, enableRightClickMenu: true, enableCloseIcon: false, genDiffToNext: () => ({
empty: true,
ownFile: "",
otherFile: "",
diff: "",
}), genDiffToPrevious: () => ({
empty: true,
ownFile: "",
otherFile: "",
diff: "",
}) }
{
selected: false, enableRightClickMenu: true, enableCloseIcon: false, genDiffToNext: () => ({
empty: true,
ownFile: "",
otherFile: "",
diff: "",
}), genDiffToPrevious: () => ({
empty: true,
ownFile: "",
otherFile: "",
diff: "",
})
}
)
const emit = defineEmits<{
'update:showMenuIdx': [v: number],
'fileItemClick': [event: MouseEvent, file: FileNodeInfo, idx: number],
@ -85,9 +90,9 @@ const taggleLikeTag = () => {
<a-dropdown :trigger="['contextmenu']" :visible="!global.longPressOpenContextMenu ? undefined : typeof idx === 'number' && showMenuIdx === idx
" @update:visible="(v: boolean) => typeof idx === 'number' && emit('update:showMenuIdx', v ? idx : -1)">
<li class="file file-item-trigger grid" :class="{
clickable: file.type === 'dir',
selected
}" :data-idx="idx" :key="file.name" draggable="true" @dragstart="emit('dragstart', $event, idx)"
clickable: file.type === 'dir',
selected
}" :data-idx="idx" :key="file.name" draggable="true" @dragstart="emit('dragstart', $event, idx)"
@dragend="emit('dragend', $event, idx)" @click.capture="emit('fileItemClick', $event, file, idx)">
<div>
@ -124,22 +129,25 @@ const taggleLikeTag = () => {
<div :key="file.fullpath" :class="`idx-${idx} item-content`" v-if="isImageFile(file.name)">
<!-- change indicators -->
<ChangeIndicator v-if="enableChangeIndicator" :gen-diff-to-next="genDiffToNext" :gen-diff-to-previous="genDiffToPrevious"/>
<ChangeIndicator v-if="enableChangeIndicator" :gen-diff-to-next="genDiffToNext"
:gen-diff-to-previous="genDiffToPrevious" />
<!-- change indicators END -->
<a-image :src="imageSrc" :fallback="fallbackImage" :preview="{
src: fullScreenPreviewImageUrl,
onVisibleChange: (v: boolean, lv: boolean) => emit('previewVisibleChange', v, lv)
}" />
src: fullScreenPreviewImageUrl,
onVisibleChange: (v: boolean, lv: boolean) => emit('previewVisibleChange', v, lv)
}" />
<div class="tags-container" v-if="customTags && cellWidth > 128">
<a-tag v-for="tag in extraTags ?? customTags" :key="tag.id" :color="tagStore.getColor(tag.name)">
{{ tag.name }}
</a-tag>
</div>
</div>
<div :class="`idx-${idx} item-content video`" v-else-if="isVideoFile(file.name)" @click="openVideoModal(file)">
<div :class="`idx-${idx} item-content video`" :urld="toVideoCoverUrl(file)" :style="{ 'background-image': `url('${toVideoCoverUrl(file)}')` }"
v-else-if="isVideoFile(file.name)" @click="openVideoModal(file)">
<div class="play-icon">
<PlayCircleFilled />
<img :src="play" style="width: 40px;height: 40px;">
</div>
<div class="tags-container" v-if="customTags && cellWidth > 128">
<a-tag v-for="tag in customTags" :key="tag.id" :color="tagStore.getColor(tag.name)">
@ -174,7 +182,6 @@ const taggleLikeTag = () => {
</a-dropdown>
</template>
<style lang="scss" scoped>
.center {
display: flex;
justify-content: center;
@ -190,6 +197,7 @@ const taggleLikeTag = () => {
overflow: hidden;
width: v-bind('$props.cellWidth + "px"');
height: v-bind('$props.cellWidth + "px"');
background-size: cover;
cursor: pointer;
}
@ -197,8 +205,9 @@ const taggleLikeTag = () => {
position: absolute;
top: 50%;
left: 50%;
font-size: 3em;
transform: translate(-50%, -50%);
transform: translate(-50%, -50%);
border-radius: 100%;
display: flex;
}
.tags-container {
@ -345,4 +354,5 @@ const taggleLikeTag = () => {
flex-direction: column;
align-items: flex-end;
}
}</style>
}
</style>

View File

@ -3,7 +3,7 @@ import { ref } from 'vue'
import * as Path from '@/util/path'
import { FileNodeInfo, mkdirs } from '@/api/files'
import { t } from '@/i18n'
import { downloadFiles, globalEvents, toRawFileUrl } from '@/util'
import { downloadFiles, globalEvents, toRawFileUrl, toStreamVideoUrl } from '@/util'
import { DownloadOutlined } from '@/icon'
import { isStandalone } from '@/util/env'
import { rebuildImageIndex } from '@/api/db'
@ -52,7 +52,7 @@ export const openVideoModal = (file: FileNodeInfo) => {
flexDirection: 'column'
}}
>
<video style={{ maxHeight: isStandalone ? '80vh' : '60vh' }} src={toRawFileUrl(file)} controls autoplay></video>
<video style={{ maxHeight: isStandalone ? '80vh' : '60vh' }} src={toStreamVideoUrl(file)} controls autoplay></video>
<div class="actions" style={{ marginTop: '16px' }}>
<Button onClick={() => downloadFiles([toRawFileUrl(file, true)])}>
{{

View File

@ -1,5 +1,7 @@
export * from '@ant-design/icons-vue'
import regex from './regex.svg'
import play from './play.svg'
export {
regex
regex,
play
}

4
vue/src/icon/play.svg Normal file
View File

@ -0,0 +1,4 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40" width="40" height="40">
<circle cx="20" cy="20" r="18" fill="#000000"/>
<polygon points="28,20 16,26 16,14" fill="#ffffff"/>
</svg>

After

Width:  |  Height:  |  Size: 195 B

View File

@ -4,14 +4,19 @@ import { uniqBy } from 'lodash-es'
const encode = encodeURIComponent
export const toRawFileUrl = (file: FileNodeInfo, download = false) =>
`${apiBase.value}/file?path=${encode(file.fullpath)}&t=${encode(file.date)}${
download ? `&disposition=${encode(file.name)}` : ''
`${apiBase.value}/file?path=${encode(file.fullpath)}&t=${encode(file.date)}${download ? `&disposition=${encode(file.name)}` : ''
}`
export const toImageThumbnailUrl = (file: FileNodeInfo, size: string = '512x512') =>
`${apiBase.value}/image-thumbnail?path=${encode(file.fullpath)}&size=${size}&t=${encode(
file.date
)}`
export const toStreamVideoUrl = (file: FileNodeInfo) =>
`${apiBase.value}/stream_video?path=${encode(file.fullpath)}`
export const toVideoCoverUrl = (file: FileNodeInfo) =>
parent.document.location.origin+ `${apiBase.value}/video_cover?path=${encode(file.fullpath)}&t=${encode(file.date)}`
export type FileTransferData = {
path: string[]
loc: string
@ -30,7 +35,7 @@ export const getFileTransferDataFromDragEvent = (e: DragEvent) => {
export const uniqueFile = (files: FileNodeInfo[]) => uniqBy(files, 'fullpath')
export function isImageFile(filename: string): boolean {
export function isImageFile (filename: string): boolean {
if (typeof filename !== 'string') {
return false
}
@ -39,7 +44,7 @@ export function isImageFile(filename: string): boolean {
return extension !== undefined && exts.includes(`.${extension}`)
}
export function isVideoFile(filename: string): boolean {
export function isVideoFile (filename: string): boolean {
if (typeof filename !== 'string') {
return false
}
@ -50,10 +55,10 @@ export function isVideoFile(filename: string): boolean {
export const isMediaFile = (file: string) => isImageFile(file) || isVideoFile(file)
export function downloadFiles(urls: string[]) {
const link = document.createElement('a');
link.style.display = 'none';
document.body.appendChild(link);
export function downloadFiles (urls: string[]) {
const link = document.createElement('a')
link.style.display = 'none'
document.body.appendChild(link)
urls.forEach((url) => {
const urlObject = new URL(url, 'https://github.com/zanllp/sd-webui-infinite-image-browsing')
@ -62,12 +67,12 @@ export function downloadFiles(urls: string[]) {
if (disposition) {
filename = disposition
}
link.href = url;
link.download = filename;
link.click();
});
link.href = url
link.download = filename
link.click()
})
document.body.removeChild(link);
document.body.removeChild(link)
}
export const downloadFileInfoJSON = (files: FileNodeInfo[], name?: string) => {