Update main.py

main.py

@@ -12,6 +12,7 @@ from nltk.tokenize import sent_tokenize
 from transformers import MarianMTModel, MarianTokenizer
 
 API_KEY = os.environ.get("API_KEY")
+VALID_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
 
 app = FastAPI()
 # CORS issue write below code
@@ -41,6 +42,11 @@ async def ocr(
     # languages: list = Body(["eng"])
 ):
     try:
+        # Check if the file format is allowed
+        file_extension = image.filename.split(".")[-1].lower()
+        if file_extension not in VALID_IMAGE_EXTENSIONS:
+            raise HTTPException(status_code=400, detail="Invalid file format. Only .jpg, .jpeg, and .png are allowed.")
+            
         content = await image.read()
         image = Image.open(BytesIO(content))
         text = pytesseract.image_to_string(image, lang = 'eng')