1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2025-01-06 03:46:47 +01:00
immich/machine_learning/app/main.py
Alex 619735fea0
Implemented image tagging using TensorFlow InceptionV3 (#28)
* Refactor docker-compose to its own folder
* Added FastAPI development environment
* Added support for GPU in docker file
* Added image classification
* creating endpoint for smart Image info
* added logo with white background on ios
* Added endpoint and trigger for image tagging
* Classify image and save into database
* Update readme
2022-02-19 22:42:10 -06:00

51 lines
1.3 KiB
Python

from typing import Optional
from pydantic import BaseModel
import numpy as np
from fastapi import FastAPI
import tensorflow as tf
from tensorflow.keras.applications import InceptionV3
from tensorflow.keras.applications.inception_v3 import preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
IMG_SIZE = 299
PREDICTION_MODEL = InceptionV3(weights='imagenet')
def warm_up():
img_path = f'./app/test.png'
img = image.load_img(img_path, target_size=(IMG_SIZE, IMG_SIZE))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)
PREDICTION_MODEL.predict(x)
# Warm up model
warm_up()
app = FastAPI()
class TagImagePayload(BaseModel):
thumbnail_path: str
@app.post("/tagImage")
async def post_root(payload: TagImagePayload):
imagePath = payload.thumbnail_path
if imagePath[0] == '.':
imagePath = imagePath[2:]
img_path = f'./app/{imagePath}'
img = image.load_img(img_path, target_size=(IMG_SIZE, IMG_SIZE))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)
preds = PREDICTION_MODEL.predict(x)
result = decode_predictions(preds, top=3)[0]
payload = []
for _, value, _ in result:
payload.append(value)
return payload