Update app/main.py

app/main.py

@@ -1,50 +1,35 @@
 import PIL
 from fastapi import FastAPI, File, UploadFile
 from pydantic import BaseModel
-from fastapi.responses import JSONResponse
 from utils.model_func import class_id_to_label, load_model, transform_image
 
-model = None
+model = None 
 app = FastAPI()
 
 
+# Create class of answer: only class name 
 class ImageClass(BaseModel):
     prediction: str
 
-class TextClass(BaseModel):
-    text: str
-
-
+# Load model at startup
 @app.on_event("startup")
-async def startup_event():
+def startup_event():
     global model
-    # Здесь используйте функцию из utils.model_func для загрузки модели
     model = load_model()
 
+@app.get('/')
+def return_info():
+    return 'Hello FastAPI'
 
-# @app.post('/classify')
-# async def classify_image(file: UploadFile = File(...)):
-#     # Здесь используйте функцию из utils.model_func для классификации изображения
-#     image_bytes = await file.read()
-#     prediction = transform_image(image_bytes, model)
-#     return {"prediction": prediction}
 
 @app.post('/classify')
-async def classify_image(file: UploadFile = File(...)):
-    # Use the function from utils.model_func to classify the image
+def classify(file: UploadFile = File(...)):
     image = PIL.Image.open(file.file)
     adapted_image = transform_image(image)
-    pred_index = model(adapted_image.unsqeeze(0).detach().cpu().numpy().argmax())
-    imagenet_class = class_id_to_label(pred_index)     
-    response = ImageClass(prediction=imagenet_class)   
-    return response
-
-          
-                       
-
-@app.post('/clf_text')
-async def classify_text(text_data: TextClass):
-    # Здесь используйте функцию из utils.model_func для классификации текста
-    text = text_data.text
-    prediction = class_id_to_label(text, model)
-    return {"prediction": prediction}
+    pred_index = model(adapted_image.unsqueeze(0)).detach().cpu().numpy().argmax()
+    imagenet_class = class_id_to_label(pred_index)
+    response = ImageClass(
+        prediction=imagenet_class
+    )
+
+    return response