README.md

@@ -9,18 +9,10 @@ In this repository, we have integrated a DeepaaS API into the  Ultralytics YOLOv
 
 # Install the API and the external submodule requirement
 To launch the API, first, install the package, and then run DeepaaS:
-``` 
+``` bash
 git clone --depth 1 https://git.scc.kit.edu/m-team/ai/yolov8_api.git
 cd  yolov8_api
 pip install -e .
-```
-
-To launch it, first install the package then run [deepaas](https://github.com/indigo-dc/DEEPaaS):
-
-```bash
-git clone https://git.scc.kit.edu/m-team/ai/yolov8_api
-cd yolov8_api
-pip install -e .
 deepaas-run --listen-ip 0.0.0.0
 ```
 

api/__init__.py

@@ -137,24 +137,28 @@ def train(**args):
             config.DATA_PATH, "raw", args["data"]
         )
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-
-        # should  the project name
         args["project"] = config.MODEL_NAME
         # point to the model directory without root directory
         args["name"] = os.path.join("models", timestamp)
-        model = YOLO(args["model"])
+        if args["weights"] is not None:
+            if os.path.isfile(args["weights"]):
+                path=args["weights"]
+            else:    
+                path= os.path.join(
+                config.MODELS_PATH, args["weights"]
+            )
+            model=YOLO(path)
+        else:
+            model = YOLO(args["model"])  
+    
         os.environ["WANDB_DISABLED"] = str(args["disable_wandb"])
-        args.pop("disable_wandb", None)
-        args.pop("task_type", None)
+        utils.pop_keys_from_dict(args, ["task_type", "disable_wandb", "weights"])
 
         model.train(**args)
         return {
             f'The model was trained successfully and was saved to: {os.path.join(args["project"], args["name"])}'
         }
 
-    except Exception as err:
-        raise HTTPException(reason=err) from err
-
 
 if __name__ == "__main__":
     fields = schemas.TrainArgsSchema().fields

api/utils.py

@@ -219,3 +219,6 @@ class DotDict:
                 setattr(self, key, DotDict(value))
             else:
                 setattr(self, key, value)
+def pop_keys_from_dict(dictionary, keys_to_pop):
+    for key in keys_to_pop:
+        dictionary.pop(key, None)