Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Upscaling and classification plugins #131

Merged
merged 16 commits into from
Jun 8, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 42 additions & 0 deletions .github/workflows/plugins-build-and-archive.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
name: Build and archive plugin solibs

on:
push:
branches: [ develop, preview ]
pull_request:
branches: [ develop, preview ]

workflow_dispatch:

jobs:
build-and-archive-plugins:
runs-on: ubuntu-latest

steps:
- uses: actions/checkout@v3

- name: Create compose stack
run: docker-compose -f docker-compose.ci.yml up -d --build --force-recreate

- name: Compile plugins
id: compile_plugins
run: docker-compose -f docker-compose.ci.yml exec -T photoprism make build-plugins

- name: Look for plugin solibs
run: docker-compose -f docker-compose.ci.yml exec -T photoprism find storage/plugins/
if: always() && steps.compile_plugins.outcome == 'success'

- name: Copy plugins from container to host
run: docker compose -f docker-compose.ci.yml cp photoprism:/go/src/github.com/photoprism/photoprism/storage/plugins/ plugins/
if: always() && steps.compile_plugins.outcome == 'success'

- name: Tear down stack
run: docker-compose -f docker-compose.ci.yml down

- name: Archive plugins
uses: actions/upload-artifact@v3
if: always() && steps.compile_plugins.outcome == 'success'
continue-on-error: true
with:
name: plugins
path: plugins/
44 changes: 44 additions & 0 deletions .github/workflows/plugins-docker.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
name: Build plugin docker images

on:
pull_request:
paths:
- "ext/**"

jobs:
publish-plugin-docker-images:
strategy:
fail-fast: false
matrix:
plugin:
- name: realesrgan
version: 0.3.0
- name: yolo8
version: 8.0.110

runs-on: ubuntu-latest
steps:
- uses: dorny/paths-filter@v2
id: changes
with:
filters: |
changed: 'ext/${{ matrix.plugin.name }}/**'

- name: Checkout
if: steps.changes.outputs.changed == 'true'
uses: actions/checkout@v3

- name: Log in to Docker Hub
if: steps.changes.outputs.changed == 'true'
uses: docker/login-action@v1
with:
username: ${{ secrets.DOCKER_USERNAME }}
password: ${{ secrets.DOCKER_PASSWORD }}

- name: Build and push Docker image
if: steps.changes.outputs.changed == 'true'
uses: docker/build-push-action@v2
with:
context: ext/${{ matrix.plugin.name }}
tags: kvalev/${{ matrix.plugin.name }}:${{ matrix.plugin.version }}
push: true
2 changes: 1 addition & 1 deletion .github/workflows/release.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ jobs:
password: ${{ secrets.DOCKER_PASSWORD }}

- name: Build and push
uses: docker/bake-action@master
uses: docker/bake-action@v3
with:
files: ./docker/docker-bake.hcl
push: true
Expand Down
16 changes: 14 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@ INSTALL_USER ?= $(DESTUID):$(DESTGID)
INSTALL_MODE ?= u+rwX,a+rX
INSTALL_MODE_BIN ?= 755

# Other parameters.
PLUGINS_PATH ?= storage/plugins

UID := $(shell id -u)
GID := $(shell id -g)
HASRICHGO := $(shell which richgo)
Expand Down Expand Up @@ -179,6 +182,8 @@ clean-local-cache:
rm -rf $(BUILD_PATH)/storage/cache/*
clean-local-config:
rm -f $(BUILD_PATH)/config/*
clean-plugins:
rm -f $(PLUGINS_PATH)/*
dep-list:
go list -u -m -json all | go-mod-outdated -direct
dep-npm:
Expand Down Expand Up @@ -219,9 +224,16 @@ build-race:
build-static:
rm -f $(BINARY_NAME)
scripts/build.sh static $(BINARY_NAME)
build-plugins: build-plugin-demo build-plugin-realesrgan build-plugin-yolo8
build-plugin-demo:
mkdir -p storage/plugins
go build -buildmode=plugin -o storage/plugins/demo.so internal/plugin/demo/demo.go
mkdir -p $(PLUGINS_PATH)
go build -tags=debug -buildmode=plugin -o $(PLUGINS_PATH)/demo.so internal/plugin/demo/demo.go
build-plugin-realesrgan:
mkdir -p $(PLUGINS_PATH)
go build -tags=debug -buildmode=plugin -o $(PLUGINS_PATH)/realesrgan.so internal/plugin/realesrgan/realesrgan.go
build-plugin-yolo8:
mkdir -p $(PLUGINS_PATH)
go build -tags=debug -buildmode=plugin -o $(PLUGINS_PATH)/yolo8.so internal/plugin/yolo8/yolo8.go
build-tensorflow:
docker build -t photoprism/tensorflow:build docker/tensorflow
docker run -ti photoprism/tensorflow:build bash
Expand Down
26 changes: 26 additions & 0 deletions docker-compose.plugins.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
version: '3.5'

## FOR TEST AND DEVELOPMENT ONLY, DO NOT USE IN PRODUCTION ##
## Setup: https://docs.photoprism.app/developer-guide/setup/ ##

services:
## PhotoPrism Development Environment
photoprism:
environment:
PHOTOPRISM_PLUGIN_YOLO8_ENABLED: "true"
PHOTOPRISM_PLUGIN_YOLO8_HOSTNAME: "yolo8"
PHOTOPRISM_PLUGIN_YOLO8_PORT: "5000"
PHOTOPRISM_PLUGIN_REALESRGAN_ENABLED: "true"
PHOTOPRISM_PLUGIN_REALESRGAN_HOSTNAME: "realesrgan"
PHOTOPRISM_PLUGIN_REALESRGAN_PORT: "5001"

## Image classification API
yolo8:
image: kvalev/yolo8:8.0.110
pull_policy: always

## Image upscaling API
realesrgan:
image: kvalev/realesrgan:0.3.0
pull_policy: always
restart: always
12 changes: 12 additions & 0 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ services:
TF_CPP_MIN_LOG_LEVEL: 0 # show TensorFlow log messages for development
## Run/install on first startup (options: update https gpu tensorflow davfs clitools clean):
PHOTOPRISM_INIT: "https tensorflow"
PHOTOPRISM_WORKERS: 1
## Hardware Video Transcoding (optional):
# PHOTOPRISM_FFMPEG_ENCODER: "nvidia" # FFmpeg encoder ("software", "intel", "nvidia", "apple", "raspberry", "vaapi") Intel: "intel" for Broadwell or later and "vaapi" for Haswell or earlier
# PHOTOPRISM_FFMPEG_ENCODER: "intel" # FFmpeg encoder ("software", "intel", "nvidia", "apple", "raspberry", "vaapi") Intel: "intel" for Broadwell or later and "vaapi" for Haswell or earlier`
Expand Down Expand Up @@ -149,6 +150,17 @@ services:
MARIADB_PASSWORD: "photoprism"
MARIADB_ROOT_PASSWORD: "photoprism"

phpmyadmin:
image: phpmyadmin/phpmyadmin:5.0
restart: always
environment:
PMA_HOST: mariadb
PMA_PORT: 4001
ports:
- 6060:80
depends_on:
- mariadb

## HTTPS Reverse Proxy (recommended) ##
## includes "*.localssl.dev" SSL certificate for test environments
## Docs: https://doc.traefik.io/traefik/
Expand Down
1 change: 1 addition & 0 deletions ext/realesrgan/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
weights
13 changes: 13 additions & 0 deletions ext/realesrgan/Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
FROM python:3.10.11-slim-buster

RUN apt-get update && apt-get install -y libgl1 libglib2.0-0

COPY requirements.txt /app/
RUN pip install -r /app/requirements.txt

COPY . /app/
WORKDIR /app

EXPOSE 5001

CMD ["python3", "main.py"]
87 changes: 87 additions & 0 deletions ext/realesrgan/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
import base64
from io import BytesIO
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from flask import Flask, request
from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
from PIL import Image
import numpy as np
import torch

import logging
import os

app = Flask(__name__)

REAL_ESRGAN_MODELS = {
"realesr-general-x4v3": {
"url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
"model_md5": "91a7644643c884ee00737db24e478156",
"scale": 4,
"model": lambda: SRVGGNetCompact(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_conv=32,
upscale=4,
act_type="prelu",
),
},
"RealESRGAN_x4plus": {
"url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
"model_md5": "99ec365d4afad750833258a1a24f44ca",
"scale": 4,
"model": lambda: RRDBNet(
num_in_ch=3,
num_out_ch=3,
num_feat=64,
num_block=23,
num_grow_ch=32,
scale=4,
),
},
}

name = "realesr-general-x4v3"

if name not in REAL_ESRGAN_MODELS:
raise ValueError(f"Unknown RealESRGAN model name: {name}")

model_info = REAL_ESRGAN_MODELS[name]

model_path = os.path.join('weights', name + '.pth')
if not os.path.isfile(model_path):
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
model_path = load_file_from_url(url=model_info["url"], model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
logging.info(f"RealESRGAN model path: {model_path}")

device = "cuda" if torch.cuda.is_available() else "cpu"

model = RealESRGANer(
scale=model_info["scale"],
model_path=model_path,
model=model_info["model"](),
half=True if "cuda" in device else False,
device=device,
)

@app.route('/superscale', methods=['POST'])
def superscale():
scale = request.json.get("scale", model_info["scale"])

image_data = base64.b64decode(request.json['image'])
image = np.asarray(Image.open(BytesIO(image_data)))
print(f"RealESRGAN input shape: {image.shape}, scale: {scale}", flush=True)

upsampled = model.enhance(image, outscale=scale)[0]
upsampled_img = Image.fromarray(upsampled)
print(f"RealESRGAN output shape: {upsampled.shape}", flush=True)

with BytesIO() as buffer:
upsampled_img.save(buffer, format="jpeg")
return {"image": base64.b64encode(buffer.getvalue()).decode()}

if __name__ == '__main__':
print("running")
app.run(host='0.0.0.0', port=5001)
2 changes: 2 additions & 0 deletions ext/realesrgan/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
realesrgan==0.3.0
flask==2.3.2
1 change: 1 addition & 0 deletions ext/yolo8/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
*.pt
13 changes: 13 additions & 0 deletions ext/yolo8/Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
FROM python:3.10.11-slim-buster

RUN apt-get update && apt-get install -y libgl1 libglib2.0-0

COPY requirements.txt /app/
RUN pip install -r /app/requirements.txt

COPY . /app/
WORKDIR /app

EXPOSE 5000

CMD ["python3", "main.py"]
43 changes: 43 additions & 0 deletions ext/yolo8/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import base64
from flask import Flask, request
from io import BytesIO
from PIL import Image
from typing import List
from ultralytics import YOLO
from ultralytics.yolo.engine.results import Results

app = Flask(__name__)

model = "yolov8n"
detect_model = YOLO(f"{model}.pt")
classify_model = YOLO(f"{model}-cls.pt")

@app.route('/hello', methods=['GET'])
def hello():
return "elloh"

@app.route('/classify', methods=['POST'])
def classify():
image_data = base64.b64decode(request.json['image'])
image = Image.open(BytesIO(image_data))

results: List[Results] = classify_model.predict(image)
result = results[0]

# take only the top3 results
take = min(len(result.names), 3)
top_n_idx = result.probs.argsort(0, descending=True)[:take].tolist()

return {result.names[idx]: result.probs[idx].item() for idx in top_n_idx}

@app.route('/detect', methods=['POST'])
def detect():
image_data = base64.b64decode(request.json['image'])
image = Image.open(BytesIO(image_data))

results: List[Results] = detect_model.predict(image)

return results[0].tojson()

if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)
2 changes: 2 additions & 0 deletions ext/yolo8/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
ultralytics==8.0.110
flask==2.3.2
4 changes: 3 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ require github.com/go-ldap/ldap/v3 v3.4.5-0.20230210083308-d16fb563008d
require (
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
github.com/creasty/defaults v1.7.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dsoprea/go-logging v0.0.0-20200710184922-b02d349568dd // indirect
github.com/dsoprea/go-utility/v2 v2.0.0-20221003172846-a3e1774ef349 // indirect
Expand All @@ -100,6 +101,7 @@ require (
github.com/leodido/go-urn v1.2.1 // indirect
github.com/mandykoh/go-parallel v0.1.0 // indirect
github.com/mattn/go-isatty v0.0.17 // indirect
github.com/mitchellh/mapstructure v1.5.0 // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/pelletier/go-toml/v2 v2.0.6 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
Expand Down Expand Up @@ -128,4 +130,4 @@ require (
golang.org/x/arch v0.2.0 // indirect
)

go 1.17
go 1.18
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDk
github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w=
github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/creasty/defaults v1.7.0 h1:eNdqZvc5B509z18lD8yc212CAqJNvfT1Jq6L8WowdBA=
github.com/creasty/defaults v1.7.0/go.mod h1:iGzKe6pbEHnpMPtfDXZEr0NVxWnPTjb1bbDy08fPzYM=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
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=
Expand Down Expand Up @@ -305,6 +307,8 @@ github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh
github.com/mattn/go-sqlite3 v1.14.0/go.mod h1:JIl7NbARA7phWnGvh0LKTyg7S9BA+6gx71ShQilpsus=
github.com/mattn/go-sqlite3 v2.0.1+incompatible h1:xQ15muvnzGBHpIpdrNi1DA5x0+TcBZzsIDwmw9uTHzw=
github.com/mattn/go-sqlite3 v2.0.1+incompatible/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
Expand Down
Loading