from typing import Optional, List
from fastapi import FastAPI, File, UploadFile, Request
from sqlmodel import Field, Session, SQLModel, create_engine, select
from pydantic import BaseModel
from network import Network
import requests
from PIL import Image as ImagePIL
import torchvision as tv

app = FastAPI()
network = Network()
network.model.eval()


class Image(SQLModel, table=True):
    key: Optional[int] = Field(default=None, primary_key=True)
    image_name: str
    label: str
    image_url: str

engine = create_engine("sqlite:///image.db")

@app.get("/")
def read_images():
    with Session(engine) as session:
        statement = select(Image)
        images = session.exec(statement).all()
        return images

class Item(BaseModel):
    key: int
    label: str = ""

@app.post("/")
def update_heroes(item: Item):
    with Session(engine) as session:
        statement = select(Image).where(Image.key == item.key)
        results = session.exec(statement)
        image = results.one()
        image.label = item.label
        session.add(image)
        session.commit()

@app.get("/predict/{item}")
def predict(item: int):
    import torch
    import numpy as np
    with Session(engine) as session:
        statement = select(Image).where(Image.key == item)
        results = session.exec(statement)
        image = results.one()
        image_url = image.image_url
        img = ImagePIL.open(requests.get(image_url, stream=True).raw)
        img = tv.transforms.functional.pil_to_tensor(img).float().unsqueeze(0)
        with torch.no_grad():
            result = network.model(img)
            result = torch.nn.functional.softmax(result)
            ret = {
                '1' : float(result[0][0]), '2': float(result[0][1])
            }
            return str(ret)


@app.post("/upload_image")
async def upload_image(files: List[UploadFile]):
    # return {"filenames": [file.filename for file in files]}
    from minio import Minio
    import io
    with Session(engine) as session:
        client = Minio(
            "localhost:9001",
            secure=False,
            access_key="<KEY>",
            secret_key="<KEY>"
        )

        for item in files:
            image_name = item.filename
            cont = await item.read()
            content = io.BytesIO(cont)
            length = len(content.read())
            content = io.BytesIO(cont)
            client.put_object(
                "image", image_name, content, length
            )
            image_url = f"http://localhost:9001/image/{image_name}"
            data = Image(image_name=image_name[:-4], label="", image_url=image_url)
            session.add(data)
            session.commit()
fastapi.FastAPI torchvision.transforms.functional.pil_to_tensor sqlmodel.create_engine minio.Minio io.BytesIO requests.get network.Network sqlmodel.select torch.no_grad sqlmodel.Session torch.nn.functional.softmax sqlmodel.Field
[((299, 308), 'fastapi.FastAPI', 'FastAPI', ([], {}), '()\n', (306, 308), False, 'from fastapi import FastAPI, File, UploadFile, Request\n'), ((319, 328), 'network.Network', 'Network', ([], {}), '()\n', (326, 328), False, 'from network import Network\n'), ((514, 549), 'sqlmodel.create_engine', 'create_engine', (['"""sqlite:///image.db"""'], {}), "('sqlite:///image.db')\n", (527, 549), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((412, 449), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (417, 449), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((593, 608), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (600, 608), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((641, 654), 'sqlmodel.select', 'select', (['Image'], {}), '(Image)\n', (647, 654), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((837, 852), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (844, 852), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1187, 1202), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1194, 1202), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1979, 1994), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1986, 1994), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((2024, 2101), 'minio.Minio', 'Minio', (['"""localhost:9001"""'], {'secure': '(False)', 'access_key': '"""<KEY>"""', 'secret_key': '"""<KEY>"""'}), "('localhost:9001', secure=False, access_key='<KEY>', secret_key='<KEY>')\n", (2029, 2101), False, 'from minio import Minio\n'), ((1544, 1559), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1557, 1559), False, 'import torch\n'), ((1622, 1657), 'torch.nn.functional.softmax', 'torch.nn.functional.softmax', (['result'], {}), '(result)\n', (1649, 1657), False, 'import torch\n'), ((2286, 2302), 'io.BytesIO', 'io.BytesIO', (['cont'], {}), '(cont)\n', (2296, 2302), False, 'import io\n'), ((2366, 2382), 'io.BytesIO', 'io.BytesIO', (['cont'], {}), '(cont)\n', (2376, 2382), False, 'import io\n'), ((885, 898), 'sqlmodel.select', 'select', (['Image'], {}), '(Image)\n', (891, 898), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1235, 1248), 'sqlmodel.select', 'select', (['Image'], {}), '(Image)\n', (1241, 1248), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1410, 1446), 'requests.get', 'requests.get', (['image_url'], {'stream': '(True)'}), '(image_url, stream=True)\n', (1422, 1446), False, 'import requests\n'), ((1466, 1509), 'torchvision.transforms.functional.pil_to_tensor', 'tv.transforms.functional.pil_to_tensor', (['img'], {}), '(img)\n', (1504, 1509), True, 'import torchvision as tv\n')]