Merge pull request #8 from xxxsen/xxxsen/feature/auto_download_dep

Xxxsen/feature/auto download dep

README.md

@@ -45,7 +45,7 @@ services:
 
 配置完成后, 使用`docker compose up` 进行刮削, 刮削完成的电影会被存储到`/data/scrape/savedir`下。
 
-**NOTE: 程序依赖go-face进行人脸识别, 以用于识别图片中的人脸并进行截图, 这个库需要有对应的模型文件, 可以通过项目目录下的`scripts/download_models.sh`进行下载, 程序会将模型下载到`models`目录, 之后将该目录移动到数据目录(`/datadir`)下即可**
+**NOTE: 程序依赖go-face进行人脸识别, 以用于识别图片中的人脸并进行截图, 这个库需要有对应的模型文件, 程序启动的时候, 会检测模型文件是否存在, 如果不存在, 则会自动下载模型文件到`数据目录`下**
 
 ### 手动编译
 

client/client.go

@@ -8,11 +8,15 @@ import (
 	"github.com/imroc/req/v3"
 )
 
+type IHTTPClient interface {
+	Do(req *http.Request) (*http.Response, error)
+}
+
 type Client struct {
 	client *http.Client
 }
 
-func NewClient() *Client {
+func NewClient() IHTTPClient {
 	// 第三方客户端用着不是很习惯, 考虑到我们需要用到的功能都是在transport里面,
 	// 所以这里直接把第三方客户端的transport提出来用...
 	reqClient := req.NewClient()
@@ -27,10 +31,6 @@ func NewClient() *Client {
 	return &Client{client: client}
 }
 
-func (c *Client) GetOriginClient() *http.Client {
-	return c.client
-}
-
 func (c *Client) Do(req *http.Request) (*http.Response, error) {
 	return c.client.Do(req)
 }

client/io.go

@@ -0,0 +1,33 @@
+package client
+
+import (
+	"compress/flate"
+	"compress/gzip"
+	"io"
+	"net/http"
+)
+
+func getResponseBody(rsp *http.Response) (io.ReadCloser, error) {
+	switch rsp.Header.Get("Content-Encoding") {
+	case "gzip":
+		return gzip.NewReader(rsp.Body)
+	case "deflate":
+		return flate.NewReader(rsp.Body), nil
+	default:
+		return rsp.Body, nil
+	}
+}
+
+func BuildReaderFromHTTPResponse(rsp *http.Response) (io.ReadCloser, error) {
+	return getResponseBody(rsp)
+}
+
+func ReadHTTPData(rsp *http.Response) ([]byte, error) {
+	defer rsp.Body.Close()
+	reader, err := BuildReaderFromHTTPResponse(rsp)
+	if err != nil {
+		return nil, err
+	}
+	defer reader.Close()
+	return io.ReadAll(reader)
+}

config/config.go

@@ -12,6 +12,11 @@ type CategoryPlugin struct {
 	Plugins []string `json:"plugins"`
 }
 
+type Dependency struct {
+	Link    string `json:"link"`
+	RelPath string `json:"rel_path"`
+}
+
 type Config struct {
 	ScanDir         string                 `json:"scan_dir"`
 	SaveDir         string                 `json:"save_dir"`
@@ -25,6 +30,7 @@ type Config struct {
 	ExtraMediaExts  []string               `json:"extra_media_exts"`
 	LogConfig       logger.LogConfig       `json:"log_config"`
 	SwitchConfig    SwitchConfig           `json:"switch_config"`
+	Dependencies    []Dependency           `json:"dependencies"`
 }
 
 type SwitchConfig struct {
@@ -62,6 +68,11 @@ func defaultConfig() *Config {
 			Level:   "info",
 			Console: true,
 		},
+		Dependencies: []Dependency{
+			{Link: "https://github.com/Kagami/go-face-testdata/raw/master/models/shape_predictor_5_face_landmarks.dat", RelPath: "models/shape_predictor_5_face_landmarks.dat"},
+			{Link: "https://github.com/Kagami/go-face-testdata/raw/master/models/dlib_face_recognition_resnet_model_v1.dat", RelPath: "models/dlib_face_recognition_resnet_model_v1.dat"},
+			{Link: "https://github.com/Kagami/go-face-testdata/raw/master/models/mmod_human_face_detector.dat", RelPath: "models/mmod_human_face_detector.dat"},
+		},
 	}
 }
 

dependency/dependency.go

@@ -0,0 +1,47 @@
+package dependency
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+	"yamdc/client"
+	"yamdc/downloadmgr"
+
+	"github.com/xxxsen/common/logutil"
+	"go.uber.org/zap"
+)
+
+const (
+	defaultSuffix = ".ts"
+)
+
+type Dependency struct {
+	URL    string
+	Target string
+}
+
+func Resolve(cli client.IHTTPClient, deps []*Dependency) error {
+	m := downloadmgr.NewManager(cli)
+	for _, dep := range deps {
+		if err := checkAndDownload(m, dep.URL, dep.Target); err != nil {
+			return fmt.Errorf("download link:%s to target:%s failed, err:%w", dep.URL, dep.Target, err)
+		}
+	}
+	return nil
+}
+
+func checkAndDownload(m *downloadmgr.DownloadManager, link string, target string) error {
+	if _, err := os.Stat(target + defaultSuffix); err == nil {
+		return nil
+	}
+	logutil.GetLogger(context.Background()).Debug("start download link", zap.String("link", link))
+	if err := m.Download(link, target); err != nil {
+		return err
+	}
+	logutil.GetLogger(context.Background()).Debug("download link succ", zap.String("link", link))
+	if err := os.WriteFile(target+defaultSuffix, []byte(fmt.Sprintf("%d", time.Now().Unix())), 0644); err != nil {
+		return fmt.Errorf("write ts file failed, err:%w", err)
+	}
+	return nil
+}

downloadmgr/.gitignore

@@ -0,0 +1 @@
+/testdata

downloadmgr/download_manager.go

@@ -0,0 +1,79 @@
+package downloadmgr
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"yamdc/client"
+)
+
+type DownloadManager struct {
+	cli client.IHTTPClient
+}
+
+func NewManager(cli client.IHTTPClient) *DownloadManager {
+	return &DownloadManager{cli: cli}
+}
+
+func (m *DownloadManager) ensureDir(dst string) error {
+	dir := filepath.Dir(dst)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("mkdir failed, path:%s, err:%w", dir, err)
+	}
+	return nil
+}
+
+func (m *DownloadManager) createHTTPStream(src string) (io.ReadCloser, error) {
+	req, err := http.NewRequest(http.MethodGet, src, nil)
+	if err != nil {
+		return nil, err
+	}
+	rsp, err := m.cli.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("do request failed, err:%w", err)
+	}
+	if rsp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("status code:%d not ok", rsp.StatusCode)
+	}
+	rc, err := client.BuildReaderFromHTTPResponse(rsp)
+	if err != nil {
+		rsp.Body.Close()
+		return nil, fmt.Errorf("build reader failed, err:%w", err)
+	}
+	return rc, nil
+}
+
+func (m *DownloadManager) writeToFile(rc io.Reader, dst string) error {
+	tmp := dst + ".temp"
+	f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
+	if err != nil {
+		return fmt.Errorf("open temp file for read failed, err:%w", err)
+	}
+
+	if _, err := io.Copy(f, rc); err != nil {
+		_ = f.Close()
+		return fmt.Errorf("transfer data failed, err:%w", err)
+	}
+	_ = f.Close()
+	if err := os.Rename(tmp, dst); err != nil {
+		return fmt.Errorf("unable to move file:%w", err)
+	}
+	return nil
+}
+
+func (m *DownloadManager) Download(src string, dst string) error {
+	if err := m.ensureDir(dst); err != nil {
+		return err
+	}
+	rc, err := m.createHTTPStream(src)
+	if err != nil {
+		return err
+	}
+	defer rc.Close()
+	if err := m.writeToFile(rc, dst); err != nil {
+		return err
+	}
+	return nil
+}

downloadmgr/download_manager_test.go

@@ -0,0 +1,14 @@
+package downloadmgr
+
+import (
+	"testing"
+	"yamdc/client"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDownloa(t *testing.T) {
+	m := NewManager(client.NewClient())
+	err := m.Download("https://github.com/Kagami/go-face-testdata/raw/master/models/shape_predictor_5_face_landmarks.dat", "testdata/abc.dat")
+	assert.NoError(t, err)
+}

main.go

@@ -8,7 +8,9 @@ import (
 	"path/filepath"
 	"strings"
 	"yamdc/capture"
+	"yamdc/client"
 	"yamdc/config"
+	"yamdc/dependency"
 	"yamdc/face"
 	"yamdc/ffmpeg"
 	"yamdc/number"
@@ -38,6 +40,9 @@ func main() {
 	if err := precheckDir(c); err != nil {
 		logkit.Fatal("precheck dir failed", zap.Error(err))
 	}
+	if err := ensureDependencies(c.DataDir, c.Dependencies); err != nil {
+		logkit.Fatal("ensure dependencies failed", zap.Error(err))
+	}
 	store.SetStorage(store.NewDiskStorage(filepath.Join(c.DataDir, "cache")))
 	if err := translator.Init(); err != nil {
 		logkit.Error("init translater failed", zap.Error(err))
@@ -165,3 +170,14 @@ func precheckDir(c *config.Config) error {
 	}
 	return nil
 }
+
+func ensureDependencies(datadir string, cdeps []config.Dependency) error {
+	deps := make([]*dependency.Dependency, 0, len(cdeps))
+	for _, item := range cdeps {
+		deps = append(deps, &dependency.Dependency{
+			URL:    item.Link,
+			Target: filepath.Join(datadir, item.RelPath),
+		})
+	}
+	return dependency.Resolve(client.NewClient(), deps)
+}

searcher/default_searcher.go

@@ -11,7 +11,6 @@ import (
 	"yamdc/model"
 	"yamdc/number"
 	"yamdc/searcher/plugin"
-	"yamdc/searcher/utils"
 	"yamdc/store"
 	"yamdc/useragent"
 
@@ -124,7 +123,7 @@ func (p *DefaultSearcher) Search(ctx context.Context, number *number.Number) (*m
 		return nil, false, fmt.Errorf("invalid http status code:%d", rsp.StatusCode)
 	}
 	defer rsp.Body.Close()
-	data, err := utils.ReadHTTPData(rsp)
+	data, err := client.ReadHTTPData(rsp)
 	if err != nil {
 		return nil, false, fmt.Errorf("read body failed, err:%w", err)
 	}
@@ -265,7 +264,7 @@ func (p *DefaultSearcher) fetchImageData(ctx *plugin.PluginContext, url string)
 	if rsp.StatusCode != http.StatusOK {
 		return nil, fmt.Errorf("get url data http code not ok, code:%d", rsp.StatusCode)
 	}
-	data, err := utils.ReadHTTPData(rsp)
+	data, err := client.ReadHTTPData(rsp)
 	if err != nil {
 		return nil, fmt.Errorf("read url data failed, err:%w", err)
 	}

searcher/utils/http_utils.go

@@ -2,37 +2,14 @@ package utils
 
 import (
 	"bytes"
-	"compress/flate"
-	"compress/gzip"
-	"io"
 	"net/http"
+	"yamdc/client"
 
 	"golang.org/x/net/html"
 )
 
-func getResponseBody(rsp *http.Response) (io.ReadCloser, error) {
-	switch rsp.Header.Get("Content-Encoding") {
-	case "gzip":
-		return gzip.NewReader(rsp.Body)
-	case "deflate":
-		return flate.NewReader(rsp.Body), nil
-	default:
-		return rsp.Body, nil
-	}
-}
-
-func ReadHTTPData(rsp *http.Response) ([]byte, error) {
-	defer rsp.Body.Close()
-	reader, err := getResponseBody(rsp)
-	if err != nil {
-		return nil, err
-	}
-	defer reader.Close()
-	return io.ReadAll(reader)
-}
-
 func ReadDataAsHTMLTree(rsp *http.Response) (*html.Node, error) {
-	data, err := ReadHTTPData(rsp)
+	data, err := client.ReadHTTPData(rsp)
 	if err != nil {
 		return nil, err
 	}