created api

app.py

@@ -1,33 +1,32 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
+import io
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    # Load model and processor from Hugging Face
-    print("Loading model and processor...")
-    processor = TrOCRProcessor.from_pretrained('tjoab/latex_finetuned')
-    model = VisionEncoderDecoderModel.from_pretrained('tjoab/latex_finetuned')
+# Load model and processor only once at startup
+processor = TrOCRProcessor.from_pretrained('tjoab/latex_finetuned')
+model = VisionEncoderDecoderModel.from_pretrained('tjoab/latex_finetuned')
 
-    # Load all images as a batch
-    sample_image = open_PIL_image("sample.png")
+@app.post("/predict")
+async def predict_latex(file: UploadFile = File(...)):
+    if file.content_type not in ["image/png", "image/jpeg"]:
+        return JSONResponse(status_code=400, content={"error": "Only PNG and JPEG files are supported."})
 
-    # Preprocess the images 
-    preproc_image = processor.image_processor(images=[sample_image], return_tensors="pt").pixel_values
+    # Read image contents
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents))
+    image = prepare_image(image)
 
-    # Generate and decode the tokens
-    # NOTE: max_length default value is very small, which often results in truncated inference if not set 
-    pred_ids = model.generate(preproc_image, max_length=128)
+    # Preprocess and run inference
+    inputs = processor(images=image, return_tensors="pt").pixel_values
+    pred_ids = model.generate(inputs, max_length=128)
     latex_preds = processor.batch_decode(pred_ids, skip_special_tokens=True)
-    return {"message": "Success", "latex_preds": latex_preds}
 
+    return {"latex_preds": latex_preds}
 
-
-# Helper funtion (path to either JPEG or PNG)
-def open_PIL_image(image_path: str) -> Image.Image:
-  image = Image.open(image_path)
-  if image_path.split('.')[-1].lower() == 'png':
-      image = Image.composite(image, PIL.Image.new('RGB', image.size, 'white'), image)
-  return image
+def prepare_image(image: Image.Image) -> Image.Image:
+    """Converts image to RGB if needed and flattens transparency if present."""
+    return Image.composite(image, Image.new('RGB', image.size, 'white'), image)

requirements.txt

@@ -2,4 +2,5 @@ fastapi
 uvicorn[standard]
 transformers
 pillow
-torch
+torch
+python-multipart