调整对接dl服务代码

Author: wwbweibo

README.md

@@ -0,0 +1,5 @@
+version: 3 
+
+services:
+  - file-server:
+

docker-compose.yaml

@@ -0,0 +1 @@
+venv/

file-server-dl/.dockerignore

@@ -0,0 +1,5 @@
+from file-server-dl:latest
+WORKDIR /app
+ADD . .
+# RUN https_proxy=http://192.168.163.65:10080 pip install -r requirements.txt
+CMD ["python", "main.py"]

file-server-dl/Dockerfile

@@ -1,16 +1,26 @@
 from magika import Magika
 from flask import Flask, request
 from service.file_understanding import FileUnderstanding
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
 
 fileUnderstanding = FileUnderstanding()
+logging.info("File Understanding Service Started")
 
 app = Flask(__name__)
 # m = Magika()
 @app.route('/api/v1/file/understanding', methods=['POST'])
 def FileTypeInterfer():
     data = request.get_json()
-    return fileUnderstanding.understand(data['path'])
+    logging.info("File Understanding Request: %s", data['path'])
+    try:
+        result = fileUnderstanding.understand(data['path'])
+        return result.to_dict()
+    except Exception as e:
+        logging.error("Error: %s", e)
+        return {"error": str(e)}
 
 if __name__ == "__main__":
     # start an  http server here
-    app.run(host='0.0.0.0', port='8081', debug=True)
+    app.run(host='0.0.0.0', port='8081', debug=False, use_reloader=False)

file-server-dl/main.py

@@ -19,5 +19,5 @@ def to_dict(self) -> dict:
             'label': self.label,
             'group': self.group,
             'description': self.description,
-            'extension': self.ext
+            'extension': self.ext.to_dict() if self.ext is not None else None
         }

file-server-dl/models/file.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from models.file import FileUnderstandingResult
 from service.image_understanding import ImageUnderstanding
+import logging
 
 class FileUnderstanding:
     def __init__(self):
@@ -13,4 +14,5 @@ def understand(self, path: str) -> FileUnderstandingResult:
         file_understanding = FileUnderstandingResult(result.output.ct_label, result.output.group, result.output.description)
         if file_understanding.group == 'image':
             file_understanding.set_ext(self.image_understanding.understand(path))
+        logging.info("File Understanding Result: %s", file_understanding)
         return file_understanding

file-server-dl/service/file_understanding.py

@@ -15,6 +15,7 @@ def __init__(self):
         self.caption_model = None
         self.caption_vis_processors = None
         self.__init_clip_model__()
+        self.__init_caption_model__()
 
     def __init_clip_model__(self):
         self.clip_model, self.clip_preprocess = load_from_name("ViT-B-16", device=self.device, download_root="./")

file-server-dl/service/image_understanding.py

@@ -0,0 +1,11 @@
+FROM golang:1.22 as builder 
+WORKDIR /build
+ADD . . 
+ENV HTTPS_PROXY=http://192.168.163.65:10080
+RUN go build -o app ./cmd/app
+
+FROM ubuntu as app
+WORKDIR /app
+COPY --from=builder /build/app .
+EXPOSE 8080
+ENTRYPOINT ["./app"]

file-server/Dockerfile

@@ -1,8 +1,11 @@
 package file
 
 import (
+	"encoding/json"
 	"fileserver/utils"
+	"strings"
 
+	"github.com/ahmetb/go-linq/v3"
 	"gorm.io/gorm"
 )
 
@@ -15,9 +18,12 @@ type File struct {
 	Group       string
 	Description string
 	Tags        string
+	Caption     string
+	Checksum    string
 }
 
 func NewFile(path string) File {
+	// calc sha256 hash checksum
 	return File{
 		Path:      path,
 		Directory: utils.GetDirectory(path),
@@ -30,3 +36,43 @@ func (f *File) SetFileType(ttype, group, description string) {
 	f.Group = group
 	f.Description = description
 }
+
+func (f *File) SetFileTypeFromUnderstanding(understanding understandingResult) {
+	f.Type = understanding.Label
+	f.Group = understanding.Group
+	f.Description = understanding.Description
+	f.setupFileExtionsionInfo(understanding)
+}
+
+func (f *File) setupFileExtionsionInfo(understanding understandingResult) {
+	if understanding.Extension == nil {
+		return
+	}
+	if understanding.Group == "image" {
+		bts, _ := json.Marshal(understanding.Extension)
+		ext := imageUnderstandingExtension{}
+		json.Unmarshal(bts, &ext)
+		f.Caption = ext.Caption
+		if len(ext.Labels) > 0 {
+			var tags []string
+			linq.From(ext.Labels).SelectT(func(lable imageUnderstandingExtensionLabel) string {
+				return lable.Label
+			}).ToSlice(&tags)
+			f.Tags = strings.Join(tags, ",")
+		}
+	}
+}
+
+func (f *File) CalcSha256() {
+	f.Checksum = utils.Sha256(f.Path)
+}
+
+type imageUnderstandingExtension struct {
+	Caption string                             `json:"caption"`
+	Labels  []imageUnderstandingExtensionLabel `json:"labels"`
+}
+
+type imageUnderstandingExtensionLabel struct {
+	Label      string `json:"label"`
+	Confidence string `json:"confidence"`
+}

file-server/file/file.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"regexp"
 	"strings"
+	"sync"
 )
 
 type ScanOptions struct {
@@ -98,77 +99,54 @@ func StartFileScanner(ctx context.Context, c chan string, repo IFileRepository)
 
 func singleFileHandler(ctx context.Context, file string, repo IFileRepository) {
 	log.Default().Printf("handling file %s", file)
-	// insert into database
-	fileType, fileGroup, fileDescription, err := FileInfer(ctx, file)
-	if err != nil {
-		log.Default().Printf("error getting file type: %v", err)
-		return
-	}
 	_file := NewFile(file)
-	_file.SetFileType(fileType, fileGroup, fileDescription)
-	// 如果为图片文件，对图像进行标注
-	if fileGroup == "image" {
-		var labels []string
-		labels, err = ImageLabel(ctx, file)
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		result, err := understanding(ctx, file)
 		if err != nil {
-			log.Default().Printf("error getting image label: %v", err)
+			log.Default().Printf("error getting file type: %v", err)
 			return
 		}
-		if len(labels) > 0 {
-			_file.Tags = strings.Join(labels, ",")
-		}
-	}
-	err = repo.CreateOrUpdateFile(ctx, _file)
+		_file.SetFileTypeFromUnderstanding(result)
+	}()
+	go func() {
+		defer wg.Done()
+		_file.Checksum = utils.Sha256(file)
+	}()
+	// insert into database
+	wg.Wait()
+	err := repo.CreateOrUpdateFile(ctx, _file)
 	if err != nil {
 		log.Default().Printf("error inserting file %s: %v", file, err)
 	}
 }
 
-// FileInfer is a function to determine the file type
-func FileInfer(ctx context.Context, file string) (t string, g string, d string, err error) {
+// understanding is a function to determine the file type, and tring to label and caption this image
+func understanding(ctx context.Context, file string) (r understandingResult, err error) {
 	data := bytes.NewBuffer([]byte(fmt.Sprintf(`{"path": "%s"}`, file)))
-	request, _ := http.NewRequest(http.MethodPost, "http://localhost:8081/api/v1/file/interfer", data)
+	request, _ := http.NewRequest(http.MethodPost, "http://192.168.163.65:8081/api/v1/file/understanding", data)
 	request.Header.Set("Content-Type", "application/json")
 	resp, err := http.DefaultClient.Do(request)
 	if err != nil {
 		log.Default().Printf("error getting file type: %v", err)
 		return
 	}
 	defer resp.Body.Close()
-	type response struct {
-		Type        string `json:"type"`
-		Group       string `json:"group"`
-		Description string `json:"description"`
-	}
+
 	bts, err := io.ReadAll(resp.Body)
 	if err != nil {
 		log.Default().Printf("error reading response body: %v", err)
 		return
 	}
-	var r response
 	json.Unmarshal(bts, &r)
-	return r.Type, r.Group, r.Description, nil
+	return
 }
 
-func ImageLabel(ctx context.Context, file string) (labels []string, err error) {
-	data := bytes.NewBuffer([]byte(fmt.Sprintf(`{"path": "%s"}`, file)))
-	request, _ := http.NewRequest(http.MethodPost, "http://localhost:8081/api/v1/file/image_label", data)
-	request.Header.Set("Content-Type", "application/json")
-	resp, err := http.DefaultClient.Do(request)
-	if err != nil {
-		log.Default().Printf("error getting file type: %v", err)
-		return
-	}
-	defer resp.Body.Close()
-	type response struct {
-		Label      string `json:"label"`
-		Confidence string `json:"confidence"`
-	}
-	var r []response
-	bts, err := io.ReadAll(resp.Body)
-	json.Unmarshal(bts, &r)
-	for _, l := range r {
-		labels = append(labels, l.Label)
-	}
-	return
+type understandingResult struct {
+	Label       string `json:"label"`
+	Group       string `json:"group"`
+	Description string `json:"description"`
+	Extension   any    `json:"extension"`
 }

file-server/file/file_scanner.go

@@ -40,8 +40,8 @@ func (r *FileRepository) ListFileByDirectory(ctx context.Context, directory stri
 
 func (r *FileRepository) CreateOrUpdateFile(ctx context.Context, file File) (err error) {
 	return r.db.Exec(`insert into files 
-	(path, directory, extension, type, "group", description, tags ) values 
-	(?,?,?,?,?,?,?) on conflict(path) do update set directory = ?, extension = ?, type = ?, "group"=?, description=? , tags = ?;`,
-		file.Path, file.Directory, file.Extension, file.Type, file.Group, file.Description, file.Tags,
-		file.Directory, file.Extension, file.Type, file.Group, file.Description, file.Tags).Error
+	(path, directory, extension, type, "group", description, tags, caption, checksum ) values 
+	(?,?,?,?,?,?,?,?,?) on conflict(path) do update set directory = ?, extension = ?, type = ?, "group"=?, description=? , tags = ?,  caption = ?, checksum = ?;`,
+		file.Path, file.Directory, file.Extension, file.Type, file.Group, file.Description, file.Tags, file.Caption, file.Checksum,
+		file.Directory, file.Extension, file.Type, file.Group, file.Description, file.Tags, file.Caption, file.Checksum).Error
 }

file-server/file/repository.go

@@ -5,9 +5,12 @@ go 1.22.4
 require (
 	go.mongodb.org/mongo-driver v1.15.1
 	gopkg.in/yaml.v2 v2.4.0
+	gorm.io/driver/sqlite v1.5.6
+	gorm.io/gorm v1.25.10
 )
 
 require (
+	github.com/ahmetb/go-linq/v3 v3.2.0
 	github.com/golang/snappy v0.0.1 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
@@ -21,6 +24,4 @@ require (
 	golang.org/x/crypto v0.17.0 // indirect
 	golang.org/x/sync v0.1.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	gorm.io/driver/sqlite v1.5.6 // indirect
-	gorm.io/gorm v1.25.10 // indirect
 )

file-server/go.mod

@@ -1,3 +1,5 @@
+github.com/ahmetb/go-linq/v3 v3.2.0 h1:BEuMfp+b59io8g5wYzNoFe9pWPalRklhlhbiU3hYZDE=
+github.com/ahmetb/go-linq/v3 v3.2.0/go.mod h1:haQ3JfOeWK8HpVxMtHHEMPVgBKiYyQ+f1/kLZh/cj9U=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=

file-server/go.sum

@@ -1,6 +1,9 @@
 package utils
 
 import (
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
 	"log"
 	"os"
 	"path"
@@ -63,3 +66,18 @@ func GetExtension(f string) string {
 	}
 	return ""
 }
+
+func Sha256(path string) string {
+	file, err := os.Open(path)
+	if err != nil {
+		log.Default().Fatalf("error opening file %s: %v", path, err)
+	}
+	defer file.Close()
+
+	hash := sha256.New()
+	if _, err := io.Copy(hash, file); err != nil {
+		log.Default().Fatalf("error calculating SHA256 for file %s: %v", path, err)
+	}
+
+	return hex.EncodeToString(hash.Sum(nil))
+}