Merge pull request #42 from johnolafenwa/johnolafenwa/custommodels

Johnolafenwa/custommodels

Author: johnolafenwa

.gitignore

@@ -136,4 +136,6 @@ dmypy.json
 .pytype/
 
 # Cython debug symbols
-cython_debug/
+cython_debug/
+
+*.out

Dockerfile.cpu

@@ -1,4 +1,4 @@
-FROM deepquestai/deepstack-base:cpu as cpu
+FROM deepquestai/deepstack-base:cpu-412378893 as cpu 
 
 ENV SLEEP_TIME 0.01
 ENV CUDA_MODE False
@@ -8,11 +8,10 @@ RUN mkdir /deeptemp
 RUN mkdir /datastore
 
 ENV DATA_DIR /datastore
-ENV TEMP_PATH /deeptemp
+ENV TEMP_PATH /deeptemp/
 ENV PROFILE desktop_cpu
 
 WORKDIR /app
-RUN pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
 RUN pip install onnxruntime==0.4.0
 
 RUN pip3 install redis

Dockerfile.gpu

@@ -1,4 +1,4 @@
-FROM deepquestai/deepstack-base:gpu as gpu
+FROM deepquestai/deepstack-base:gpu-412378893 as gpu
 
 ENV SLEEP_TIME 0.01
 ENV CUDA_MODE True
@@ -8,11 +8,10 @@ RUN mkdir /deeptemp
 RUN mkdir /datastore
 
 ENV DATA_DIR /datastore
-ENV TEMP_PATH /deeptemp
+ENV TEMP_PATH /deeptemp/
 ENV PROFILE desktop_gpu
 
 WORKDIR /app
-RUN pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
 RUN pip install onnxruntime==0.5.0
 
 RUN pip3 install redis

deepstack/intelligencelayer/shared/detection.py

@@ -5,7 +5,7 @@
 import io
 import _thread as thread
 from multiprocessing import Process
-from PIL import Image
+from PIL import Image,UnidentifiedImageError
 import torch.nn.functional as F
 import ast
 import sqlite3
@@ -21,7 +21,13 @@
 import torchvision.transforms as transforms
 import traceback
 from PIL import UnidentifiedImageError
+import argparse
 
+parser = argparse.ArgumentParser()
+parser.add_argument("--model",type=str,default=None)
+parser.add_argument("--name",type=str,default=None)
+
+opt = parser.parse_args()
 
 def objectdetection(thread_name: str, delay: float):
 
@@ -31,9 +37,15 @@ def objectdetection(thread_name: str, delay: float):
     db = SharedOptions.db
     TEMP_PATH = SharedOptions.TEMP_PATH
 
-    IMAGE_QUEUE = "detection_queue"
+    if opt.name == None:
+        IMAGE_QUEUE = "detection_queue"
+    else:
+        IMAGE_QUEUE = opt.name+"_queue"
 
-    model_name  = SharedOptions.SETTINGS.DETECTION_MODEL
+    if opt.model == None:
+        model_path  = os.path.join(SHARED_APP_DIR,SharedOptions.SETTINGS.DETECTION_MODEL)
+    else:
+        model_path = opt.model
 
     if MODE == "High":
 
@@ -47,7 +59,7 @@ def objectdetection(thread_name: str, delay: float):
 
         reso = SharedOptions.SETTINGS.DETECTION_LOW
 
-    detector = YOLODetector(os.path.join(SHARED_APP_DIR,model_name),reso,cuda=CUDA_MODE)
+    detector = YOLODetector(model_path,reso,cuda=CUDA_MODE)
     while True:
         queue = db.lrange(IMAGE_QUEUE,0,0)
 
@@ -86,8 +98,8 @@ def objectdetection(thread_name: str, delay: float):
 
                         outputs.append(detection)
 
-                    response = {"success":True,"predictions":outputs}
-                    
+                    output = {"success":True,"predictions":outputs}
+
                 except UnidentifiedImageError:
                     err_trace = traceback.format_exc()
                     print(err_trace,file=sys.stderr,flush=True)

deepstack/intelligencelayer/shared/face.py

@@ -388,8 +388,8 @@ def face(thread_name,delay):
 
                         output = {"success":False, "error":"error occured on the server","code":500}
 
-                     finally:
-                         SharedOptions.db.set(req_id,json.dumps(output))
+                    finally:
+                        SharedOptions.db.set(req_id,json.dumps(output))
 
                         if os.path.exists(SharedOptions.TEMP_PATH+img_id):
                             os.remove(SharedOptions.TEMP_PATH+img_id)

deepstack/intelligencelayer/shared/scene.py

@@ -5,7 +5,7 @@
 import io
 import _thread as thread
 from multiprocessing import Process
-from PIL import Image
+from PIL import Image,UnidentifiedImageError
 import torch.nn.functional as F
 import ast
 import sqlite3