upload
ihmily
11 months ago
parent
f117352b3e
commit
c53ad69b39
@ -0,0 +1,119 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import sys
|
||||
from fastapi import FastAPI, File, UploadFile, Form, Response
|
||||
from fastapi import Request
|
||||
import requests
|
||||
import cv2
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.utils.constant import Tasks
|
||||
from modelscope.outputs import OutputKeys
|
||||
import numpy as np
|
||||
from starlette.staticfiles import StaticFiles
|
||||
from starlette.templating import Jinja2Templates
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
model_paths = {
|
||||
"universal": {'path': 'damo/cv_unet_universal-matting', 'task': Tasks.universal_matting},
|
||||
"people": {'path': 'damo/cv_unet_image-matting', 'task': Tasks.portrait_matting},
|
||||
}
|
||||
|
||||
default_model = list(model_paths.keys())[0]
|
||||
default_model_info = model_paths[default_model]
|
||||
loaded_models = {default_model: pipeline(default_model_info['task'], model=default_model_info['path'])}
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
def __init__(self):
|
||||
self.loaded_models = {default_model: loaded_models[default_model]}
|
||||
|
||||
def load_model(self, model_name):
|
||||
if model_name not in self.loaded_models:
|
||||
model_info = model_paths[model_name]
|
||||
model_path = model_info['path']
|
||||
task_group = model_info['task']
|
||||
|
||||
self.loaded_models[model_name] = pipeline(task_group, model=model_path)
|
||||
return self.loaded_models[model_name]
|
||||
|
||||
|
||||
model_loader = ModelLoader()
|
||||
|
||||
|
||||
@app.post("/switch_model/{new_model}")
|
||||
async def switch_model(new_model: str):
|
||||
if new_model not in model_paths:
|
||||
return {"content": "Invalid model selection"}, 400
|
||||
model_info = model_paths[new_model]
|
||||
|
||||
loaded_models[new_model] = pipeline(model_info['task'], model=model_info['path'])
|
||||
model_loader.loaded_models = loaded_models
|
||||
return {"content": f"Switched to model: {new_model}"}, 200
|
||||
|
||||
|
||||
@app.post("/matting")
|
||||
async def matting(image: UploadFile = File(...), model: str = Form(default=default_model, alias="model")):
|
||||
image_bytes = await image.read()
|
||||
img = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
|
||||
|
||||
if model not in model_paths:
|
||||
return {"content": "Invalid model selection"}, 400
|
||||
|
||||
selected_model = model_loader.load_model(model)
|
||||
|
||||
result = selected_model(img)
|
||||
|
||||
output_img = result[OutputKeys.OUTPUT_IMG]
|
||||
|
||||
output_bytes = cv2.imencode('.png', output_img)[1].tobytes()
|
||||
|
||||
return Response(content=output_bytes, media_type='image/png')
|
||||
|
||||
|
||||
@app.post("/matting/url")
|
||||
async def matting_url(request: Request, model: str = Form(default=default_model, alias="model")):
|
||||
try:
|
||||
json_data = await request.json()
|
||||
image_url = json_data.get("image_url")
|
||||
except Exception as e:
|
||||
return {"content": f"Error parsing JSON data: {str(e)}"}, 400
|
||||
|
||||
if not image_url:
|
||||
return {"content": "Image URL is required"}, 400
|
||||
|
||||
response = requests.get(image_url)
|
||||
if response.status_code != 200:
|
||||
return {"content": "Failed to fetch image from URL"}, 400
|
||||
|
||||
img_array = np.frombuffer(response.content, dtype=np.uint8)
|
||||
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
||||
|
||||
if model not in model_paths:
|
||||
return {"content": "Invalid model selection"}, 400
|
||||
|
||||
selected_model = model_loader.load_model(model)
|
||||
|
||||
result = selected_model(img)
|
||||
|
||||
output_img = result[OutputKeys.OUTPUT_IMG]
|
||||
|
||||
output_bytes = cv2.imencode('.png', output_img)[1].tobytes()
|
||||
|
||||
return Response(content=output_bytes, media_type='image/png')
|
||||
|
||||
|
||||
templates = Jinja2Templates(directory="web")
|
||||
app.mount("/static", StaticFiles(directory="web/static"), name="static")
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def read_index(request: Request):
|
||||
return templates.TemplateResponse("index.html", {"request": request, "default_model": default_model,
|
||||
"available_models": list(model_paths.keys())})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
defult_bind_host = "0.0.0.0" if sys.platform != "win32" else "127.0.0.1"
|
||||
uvicorn.run(app, host=defult_bind_host, port=8000)
|
||||
@ -0,0 +1,10 @@
|
||||
fastapi==0.108.0
|
||||
opencv-python
|
||||
requests
|
||||
modelscope>=1.10.0
|
||||
torch==2.1.2
|
||||
transformers==4.36.2
|
||||
sentencepiece==0.1.*
|
||||
tensorflow
|
||||
python-multipart
|
||||
uvicorn
|
||||
@ -0,0 +1,145 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Image Matting</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: 'Arial', sans-serif;
|
||||
background-color: #f4f4f4;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
header {
|
||||
background-color: #333;
|
||||
color: white;
|
||||
padding: 1em;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
main {
|
||||
max-width: 800px;
|
||||
margin: 2em auto;
|
||||
background-color: white;
|
||||
padding: 2em;
|
||||
box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
#upload-form {
|
||||
text-align: center;
|
||||
margin-bottom: 2em;
|
||||
}
|
||||
|
||||
#images-container {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
#original-img,
|
||||
#result-img {
|
||||
max-width: 45%;
|
||||
height: auto;
|
||||
border: 1px solid #ccc;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
|
||||
<body>
|
||||
<header>
|
||||
<h1>Image Matting</h1>
|
||||
</header>
|
||||
<a href="https://github.com/ihmily"><img style="position: absolute; top: 0; right: 0; border: 0;" decoding="async" width="149" height="149" src="/static/images/forkme_right_gray_6d6d6d.png" class="attachment-full size-full" alt="Fork me on GitHub" loading="lazy" data-recalc-dims="1"></a>
|
||||
<main>
|
||||
<form id="upload-form" enctype="multipart/form-data">
|
||||
<label for="image-upload">Select an image:</label>
|
||||
<input type="file" id="image-upload" accept="image/*" required>
|
||||
|
||||
<label for="model-select">Select a model:</label>
|
||||
<select id="model-select" name="model">
|
||||
{% for model in available_models %}
|
||||
<option value="{{ model }}" {% if model == default_model %}selected{% endif %}>{{ model }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
|
||||
<button type="button" id="run" onclick="uploadImage()">Upload</button>
|
||||
</form>
|
||||
|
||||
<div id="images-container">
|
||||
<img id="original-img" alt=" " src="data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7">
|
||||
<img id="result-img" alt=" " src="data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7">
|
||||
</div>
|
||||
</main>
|
||||
|
||||
<script>
|
||||
document.getElementById("image-upload").addEventListener("change", function() {
|
||||
const inputElement = this;
|
||||
const file = inputElement.files[0];
|
||||
|
||||
if (file) {
|
||||
const originalImgElement = document.getElementById("original-img");
|
||||
originalImgElement.src = URL.createObjectURL(file);
|
||||
document.getElementById("result-img").src = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7";
|
||||
}
|
||||
});
|
||||
|
||||
async function uploadImage() {
|
||||
const inputElement = document.getElementById("image-upload");
|
||||
const file = inputElement.files[0];
|
||||
|
||||
if (file) {
|
||||
const formData = new FormData();
|
||||
formData.append("image", file);
|
||||
|
||||
const modelSelect = document.getElementById("model-select");
|
||||
const selectedModel = modelSelect.value;
|
||||
|
||||
formData.append("model", selectedModel);
|
||||
|
||||
const response = await fetch("/matting", {
|
||||
method: "POST",
|
||||
body: formData,
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const resultBlob = await response.blob();
|
||||
const resultImgElement = document.getElementById("result-img");
|
||||
resultImgElement.src = URL.createObjectURL(resultBlob);
|
||||
} else {
|
||||
alert("Matting failed. Please try again.");
|
||||
}
|
||||
} else {
|
||||
alert("Please select an image.");
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById("model-select").addEventListener("change", function() {
|
||||
const newModel = this.value;
|
||||
|
||||
alert(`Switching to model: ${newModel}`);
|
||||
|
||||
fetch(`/switch_model/${newModel}`, {
|
||||
method: "POST",
|
||||
})
|
||||
.then(response => {
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
return response.json();
|
||||
})
|
||||
.then(data => {
|
||||
alert(`Switched to model: ${newModel}`);
|
||||
})
|
||||
.catch(error => {
|
||||
alert(`Model switch failed. Please try again. Error: ${error}`);
|
||||
});
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
|
||||
</html>
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 3.7 KiB |
Loading…
Reference in New Issue