Update main.py

main.py

@@ -3,6 +3,8 @@ import os
 from hiou import hugging
 import json
 import io
+import random
+import string
 
 def get_answer(image: io.BytesIO, question: str):
     filename = image.name
@@ -80,12 +82,24 @@ app = Flask(__name__)
 @app.route('/getans', methods=['POST'])
 def gettans_task():
     if 'file' not in request.files:
-        return jsonify({"status" : False, "msg" : "No file found"}), 400
-    file = request.files['file']
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
-    file_content = io.BytesIO(file.read())
-    file_content.name = file.filename
+        if len(request.data) > 0:
+            pass
+        else:
+            return jsonify({"status" : False, "msg" : "No file found"}), 400
+
+    if 'file' in request.files:
+        file = request.files['file']
+        if file.filename == '':
+            return jsonify({"error": "No selected file"}), 400
+        file_content = io.BytesIO(file.read())
+        file_content.name = file.filename
+    else:
+        rawFile = request.data
+        file = io.BytesIO(rawFile)
+        if not request.args.get("name"):
+            filename = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(10)) + ".png"
+        file.name = request.args.get("name", filename)
+
     headers = dict(request.headers)
     if not headers.get("KEY") != os.getenv("KEY"):
         return jsonify({"status" : False, "msg" : "Invalid API Key"}), 404
@@ -96,12 +110,24 @@ def gettans_task():
 @app.route('/ocr', methods=['POST'])
 def task_ocr():
     if 'file' not in request.files:
-        return jsonify({"status" : False, "msg" : "No file found", "data" : f"{request.files}"}), 400
-    file = request.files['file']
-    if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
-    file_content = io.BytesIO(file.read())
-    file_content.name = file.filename
+        if len(request.data) > 0:
+            pass
+        else:
+            return jsonify({"status" : False, "msg" : "No file found"}), 400
+
+    if 'file' in request.files:
+        file = request.files['file']
+        if file.filename == '':
+            return jsonify({"error": "No selected file"}), 400
+        file_content = io.BytesIO(file.read())
+        file_content.name = file.filename
+    else:
+        rawFile = request.data
+        file = io.BytesIO(rawFile)
+        if not request.args.get("name"):
+            filename = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(10)) + ".png"
+        file.name = request.args.get("name", filename)
+
     headers = dict(request.headers)
     if not headers.get("KEY") != os.getenv("KEY"):
         return jsonify({"status" : False, "msg" : "Invalid API Key"}), 404