添加视频帧提取功能和阿里云OSS存储支持

- 新增从视频素材提取首帧/尾帧的功能,支持画面连续性编辑
- 添加阿里云OSS存储支持,可配置本地或OSS存储方式
- 导入视频素材时自动探测并更新视频时长信息
- 前端添加从素材提取尾帧的UI界面
- 添加FramePrompt模型的数据库迁移

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
empty
2026-01-18 21:44:39 +08:00
parent fe595db96e
commit d970107a34
13 changed files with 351 additions and 19 deletions

View File

@@ -2,27 +2,48 @@ package services
import (
"fmt"
"os"
"path/filepath"
"strconv"
"strings"
"time"
models "github.com/drama-generator/backend/domain/models"
"github.com/drama-generator/backend/infrastructure/external/ffmpeg"
"github.com/drama-generator/backend/infrastructure/storage"
"github.com/drama-generator/backend/pkg/config"
"github.com/drama-generator/backend/pkg/logger"
"gorm.io/gorm"
)
type AssetService struct {
db *gorm.DB
log *logger.Logger
ffmpeg *ffmpeg.FFmpeg
db *gorm.DB
log *logger.Logger
ffmpeg *ffmpeg.FFmpeg
cfg *config.Config
ossStorage *storage.OssStorage
}
func NewAssetService(db *gorm.DB, log *logger.Logger) *AssetService {
return &AssetService{
func NewAssetService(db *gorm.DB, log *logger.Logger, cfg ...*config.Config) *AssetService {
service := &AssetService{
db: db,
log: log,
ffmpeg: ffmpeg.NewFFmpeg(log),
}
if len(cfg) > 0 {
service.cfg = cfg[0]
// 如果配置了 OSS初始化 OSS 存储
if cfg[0].Storage.Type == "oss" && storage.IsOssConfigured(&cfg[0].Storage.Oss) {
ossStorage, err := storage.NewOssStorage(&cfg[0].Storage.Oss)
if err != nil {
log.Warnw("Failed to initialize OSS storage, falling back to local", "error", err)
} else {
service.ossStorage = ossStorage
log.Infow("OSS storage initialized", "bucket", cfg[0].Storage.Oss.BucketName)
}
}
}
return service
}
type CreateAssetRequest struct {
@@ -264,6 +285,23 @@ func (s *AssetService) ImportFromVideoGen(videoGenID uint) (*models.Asset, error
storyboardNum = &videoGen.Storyboard.StoryboardNumber
}
// 如果 duration 为空,尝试使用 FFmpeg 探测
duration := videoGen.Duration
if duration == nil || *duration == 0 {
s.log.Infow("Duration is empty, probing video duration", "video_gen_id", videoGenID)
probedDuration, err := s.ffmpeg.GetVideoDuration(*videoGen.VideoURL)
if err == nil && probedDuration > 0 {
durationInt := int(probedDuration + 0.5) // 四舍五入
duration = &durationInt
s.log.Infow("Probed video duration", "video_gen_id", videoGenID, "duration", durationInt)
// 同时更新 VideoGeneration 表
s.db.Model(&videoGen).Update("duration", durationInt)
} else {
s.log.Warnw("Failed to probe video duration", "video_gen_id", videoGenID, "error", err)
}
}
asset := &models.Asset{
Name: fmt.Sprintf("Video_%d", videoGen.ID),
Type: models.AssetTypeVideo,
@@ -273,7 +311,7 @@ func (s *AssetService) ImportFromVideoGen(videoGenID uint) (*models.Asset, error
StoryboardID: videoGen.StoryboardID,
StoryboardNum: storyboardNum,
VideoGenID: &videoGenID,
Duration: videoGen.Duration,
Duration: duration,
Width: videoGen.Width,
Height: videoGen.Height,
}
@@ -283,8 +321,125 @@ func (s *AssetService) ImportFromVideoGen(videoGenID uint) (*models.Asset, error
}
if err := s.db.Create(asset).Error; err != nil {
return nil, fmt.Errorf("failed to create asset: %w", err)
return nil, fmt.Errorf("failed to create asset: %w\n", err)
}
return asset, nil
}
// ExtractFrameRequest 视频帧提取请求
type ExtractFrameRequest struct {
Position string `json:"position"` // "first" 或 "last"
StoryboardID uint `json:"storyboard_id"` // 关联的分镜 ID
FrameType string `json:"frame_type"` // image_generations 的帧类型,默认 "first"
}
// ExtractFrameFromAsset 从视频素材中提取帧并保存到 image_generations 表
func (s *AssetService) ExtractFrameFromAsset(assetID uint, req *ExtractFrameRequest) (*models.ImageGeneration, error) {
// 获取素材
var asset models.Asset
if err := s.db.First(&asset, assetID).Error; err != nil {
return nil, fmt.Errorf("asset not found: %w", err)
}
// 验证是否为视频素材
if asset.Type != models.AssetTypeVideo {
return nil, fmt.Errorf("asset is not a video")
}
// 默认值
position := req.Position
if position == "" {
position = "last"
}
frameType := req.FrameType
if frameType == "" {
frameType = "first" // 提取的尾帧用作下一个分镜的首帧
}
// 生成唯一文件名
timestamp := time.Now().Format("20060102_150405")
fileName := fmt.Sprintf("frame_%s_%d_%s.jpg", position, assetID, timestamp)
// 创建临时目录
tmpDir := "./data/tmp"
if err := os.MkdirAll(tmpDir, 0755); err != nil {
return nil, fmt.Errorf("failed to create temp directory: %w", err)
}
outputPath := filepath.Join(tmpDir, fileName)
// 使用 FFmpeg 提取帧
_, err := s.ffmpeg.ExtractFrame(asset.URL, outputPath, position)
if err != nil {
return nil, fmt.Errorf("failed to extract frame: %w", err)
}
var imageURL string
// 根据存储类型上传
if s.ossStorage != nil {
// 上传到 OSS
url, err := s.ossStorage.UploadWithFilename(outputPath, "extracted_frames", fileName)
if err != nil {
s.log.Errorw("Failed to upload to OSS, falling back to local", "error", err)
} else {
imageURL = url
s.log.Infow("Frame uploaded to OSS", "url", url)
// 删除临时文件
os.Remove(outputPath)
}
}
// 如果 OSS 上传失败或未配置,使用本地存储
if imageURL == "" {
localPath := "./data/storage"
baseURL := "/static"
if s.cfg != nil && s.cfg.Storage.LocalPath != "" {
localPath = s.cfg.Storage.LocalPath
}
if s.cfg != nil && s.cfg.Storage.BaseURL != "" {
baseURL = s.cfg.Storage.BaseURL
}
// 移动文件到最终位置
finalDir := filepath.Join(localPath, "extracted_frames")
if err := os.MkdirAll(finalDir, 0755); err != nil {
return nil, fmt.Errorf("failed to create output directory: %w", err)
}
finalPath := filepath.Join(finalDir, fileName)
if err := os.Rename(outputPath, finalPath); err != nil {
return nil, fmt.Errorf("failed to move file: %w", err)
}
imageURL = fmt.Sprintf("%s/extracted_frames/%s", baseURL, fileName)
}
// 获取 DramaID
var dramaID uint
if asset.DramaID != nil {
dramaID = *asset.DramaID
}
// 创建 image_generation 记录
imageGen := &models.ImageGeneration{
DramaID: dramaID,
StoryboardID: &req.StoryboardID,
Prompt: fmt.Sprintf("Extracted %s frame from video asset #%d", position, assetID),
Status: models.ImageStatusCompleted,
ImageURL: &imageURL,
FrameType: &frameType,
}
if err := s.db.Create(imageGen).Error; err != nil {
return nil, fmt.Errorf("failed to create image generation record: %w", err)
}
s.log.Infow("Frame extracted and saved",
"asset_id", assetID,
"position", position,
"image_gen_id", imageGen.ID,
"frame_type", frameType,
"image_url", imageURL)
return imageGen, nil
}