fix: Update for --files instead of --input, make accept option with d…

…efault application/json
ai4os-hub · Jul 12, 2024 · 254651d · 254651d
1 parent dd7c6ee
commit 254651d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/api/__init__.py b/api/__init__.py
@@ -268,12 +268,12 @@ def main():
             logger.debug("Calling method with args: %s", args)
             del vars(args)["method"]
             if hasattr(args, "files"):
-                file_extension = os.path.splitext(args.input)[1]
+                file_extension = os.path.splitext(args.files)[1]
                 args.files = UploadedFile(
-                    "input",
+                    "files",
                     args.files,
                     "application/octet-stream",
-                    f"input{file_extension}",
+                    f"files{file_extension}",
                 )
             results = method_function(**vars(args))
         print(json.dumps(results))
@@ -322,7 +322,7 @@ def main():
     --data /srv/football-players-detection-7/data.yaml\
     --Enable_MLFLOW --epochs 50
     
-        python3 api/__init__.py  predict --input \
+        python3 api/__init__.py  predict --files \
     /srv/yolov8_api/tests/data/det/test/cat1.jpg\
     --task_type  det --accept application/json
     """
diff --git a/api/schemas.py b/api/schemas.py
@@ -97,6 +97,7 @@ class Meta:
             "description": "image size as scalar or (h, w) list,"
             " i.e. (640, 480)"
         },
+        load_default=[640,480]
     )
 
     conf = fields.Float(
@@ -153,7 +154,8 @@ class Meta:
             "description": "Return format for method response.",
             "location": "headers",
         },
-        required=True,
+        required=False,
+        load_default="application/json",
         validate=validate.OneOf(responses.content_types),
     )
 

diff --git a/yolov8_api/__init__.py b/yolov8_api/__init__.py
@@ -47,7 +47,7 @@ def predict(
     print("arg of prediction are", args)
 
     model = YOLO(args["model"])
-    test_image_path = args["input"]
+    test_image_path = args["files"]
     results = []
     for image_path in test_image_path:
         print(
@@ -57,7 +57,7 @@ def predict(
         utils.remove_keys_from_dict(
             args,
             [
-                "input",
+                "files",
                 "accept",
                 "task_type",
             ],
@@ -73,7 +73,7 @@ def predict(
 
 if __name__ == "__main__":
     args = {
-        "input": ["/home/se1131/cat1.jpg"],
+        "files": ["/home/se1131/cat1.jpg"],
         "model": "yolov8n.pt",
         "imgsz": [680, 512],
         "conf": 0.25,