media_enhancer / app.py
d0tpy's picture
Update app.py
d03c20e verified
raw
history blame
2.32 kB
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from ./image_enhancer import EnhancementMethod, Enhancer
from vidoe_enhancer import VideoEnhancer
from pydantic import BaseModel
from PIL import Image
from io import BytesIO
import base64
import numpy as np
import magic
class EnhancementRequest(BaseModel):
method: EnhancementMethod = EnhancementMethod.gfpgan
background_enhancement: bool = True
upscale: int = 2
class _EnhanceBase(BaseModel):
encoded_base_img: List[str]
app = FastAPI()
@app.get("/")
def greet_json():
return {"Initializing GlamApp Enhancer"}
@app.post("/enhance/image/")
async def enhance_image(
file: UploadFile = File(...),
request: EnhancementRequest = EnhancementRequest()
):
try:
if not file.content_type.startswith('image/'):
raise HTTPException(status_code=400, detail="Invalid file type")
contents = await file.read()
base64_encoded_image = base64.b64encode(contents).decode('utf-8')
data = _EnhanceBase(encoded_base_img=[base64_encoded_image])
enhancer = Enhancer(request.method, request.background_enhancement, request.upscale)
enhanced_img, original_resolution, enhanced_resolution = await enhancer.enhance(data)
enhanced_image = Image.fromarray(enhanced_img)
img_byte_arr = BytesIO()
enhanced_image.save(img_byte_arr, format='PNG')
img_byte_arr.seek(0)
return StreamingResponse(img_byte_arr, media_type="image/png")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/enhance/video/")
async def enhance_video(file: UploadFile = File(...)):
enhancer = VideoEnhancer()
file_header = await file.read(1024)
file.file.seek(0)
mime = magic.Magic(mime=True)
file_mime_type = mime.from_buffer(file_header)
accepted_mime_types = [
'video/mp4',
'video/mpeg',
'video/x-msvideo',
'video/quicktime',
'video/x-matroska',
'video/webm'
]
if file_mime_type not in accepted_mime_types:
raise HTTPException(status_code=400, detail="Invalid file type. Please upload a video file.")
return await enhancer.stream_enhanced_video(file)