-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
60 lines (52 loc) · 1.79 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
from fastapi import FastAPI, File, UploadFile, Form,HTTPException
from pydantic import BaseModel
from typing import List
import io
import numpy as np
import sys
import logging
from model import Prediction
from Palette import palettedImage
import json
import os
import pandas as pd
from PIL import Image
import logging
logger = logging.getLogger('dev')
logger.setLevel(logging.INFO)
fileHandler = logging.FileHandler('test.log')
fileHandler.setLevel(logging.DEBUG)
formatter = logging.Formatter('[QUERY] %(asctime)s - %(image_name)s - %(srv)s - %(message)s')
fileHandler.setFormatter(formatter)
logger.addHandler(fileHandler)
app = FastAPI()
@app.get('/')
def root_route():
return { 'error': 'Use POST /prediction instead of the root route!' }
@app.post('/prediction/', response_model=Prediction)
async def prediction_route(file: UploadFile = File(...), srv : str = Form(...)):
if file.content_type.startswith('image/') is False:
extr = {'srv': srv, 'image_name': file.filename}
logger.error("ERROR", extra=extr)
raise HTTPException(status_code=400, detail=f'File \'{file.filename}\' is not an image.')
try:
contents = await file.read()
pil_image = Image.open(io.BytesIO(contents))
if pil_image.mode == 'RGBA':
pil_image = pil_image.convert('RGB')
pim = palettedImage(pil_image,file.filename,clusters=6, colorOffset=3)
pim.palettize()
#pim.save_paletted()
data = pim.get_params()
extr = {'srv' : srv, 'image_name' : file.filename}
logger.info("OK", extra=extr)
return {
'filename' : file.filename,
'contenttype': file.content_type,
'prediction' : data
}
except:
e = sys.exc_info()[1]
extr = {'srv': srv, 'image_name': None}
logger.error(f"ERROR: {str(e)}", extra=extr)
raise HTTPException(status_code=500, detail=str(e))