YACWC

utils.py

@@ -38,7 +38,7 @@ r = redis.Redis('localhost',port=6379, db=14)
 
 logger = logging.getLogger('live_inference')
 
-def get_snap( url, username, password, proxies = None, timeout=5, save_image = None, camera_name = 'N/A'):
+def get_snap( url, username, password, proxies = None, timeout=5, save_image = None, camera_name = 'N/A', width = None, height = None):
     data = {
         'cmd': 'Snap',
         'channel': 0,
@@ -47,6 +47,9 @@ def get_snap( url, username, password, proxies = None, timeout=5, save_image = N
         'user': username,
         'password': password,
     }
+    if width is not None and height is not None:
+        data['width'] = width
+        data['height'] = height
     parms = parse.urlencode(data, safe="!").encode("utf-8")
 
     try:
@@ -62,10 +65,11 @@ def get_snap( url, username, password, proxies = None, timeout=5, save_image = N
 
             logging.info(f'{camera_name}: Got image of {img_rgb.shape}')
             return img_rgb
-
+        else:
+            logging.info(f'{camera_name}: Got response code of {response.status_code}')
     except Exception as e:
         logging.error(f'{camera_name} failure: {str(e)}')
-        raise
+
 
 
 # %%
@@ -128,8 +132,9 @@ def dump_model_results_to_json(camera_name, timestamp,  output_array, hash_value
 
     json_str = json.dumps(round_floats(score_dict))
 
-
-    with open('/home/thebears/source/infer/scores/' + camera_name,'a') as ff:
+    rt_path = '/home/thebears/source/infer/scores/'
+    os.makedirs(rt_path, exist_ok = True)
+    with open(rt_path  + camera_name,'a') as ff:
         ff.write(json_str)
         ff.write('\n')
 
@@ -177,7 +182,7 @@ def run_model(img_scoring_queue):
                     continue
 
 class SnapManager():
-    def __init__(self, ip, url_api, username, password, camera_name, msg_queue=None, img_scoring_queue=None, split_into_two=False, **kwargs):
+    def __init__(self, ip, url_api, username, password, camera_name, msg_queue=None, img_scoring_queue=None, split_into_two=False, resolution = None, **kwargs):
         self.ip = ip
         self.url_api = url_api
         self.username = username
@@ -186,6 +191,7 @@ class SnapManager():
         self.split_into_two = split_into_two
         self.msg_queue = msg_queue
         self.img_scoring_queue = img_scoring_queue
+        self.resolution = resolution
         logger.info(f"{self.camera_name}: initialized")
 
     def format_image_for_model(self, image, timestamp):
@@ -217,8 +223,14 @@ class SnapManager():
             })
         return msg
 
-    def capture_and_prepare(self):
-        img = get_snap(self.url_api,self.username, self.password, camera_name =  self.camera_name)
+    def capture_and_prepare(self, save_image = None):
+        if self.resolution is not None:
+            width = self.resolution[1]
+            height = self.resolution[0]
+        else:
+            width = None
+            height = None
+        img = get_snap(self.url_api,self.username, self.password, camera_name =  self.camera_name, save_image= save_image, width = width, height = height)
         if img is not None:
             timestamp = time.time()
             return self.format_image_for_model(img, timestamp)
@@ -230,9 +242,13 @@ class SnapManager():
                 msg = self.msg_queue.get(timeout=0.1)
                 if msg == 'exit':
                     break
-                if msg == 'get':
+                elif msg.startswith('get'):
+                    if '+save' in msg:
+                        save_image = 'images/'+self.camera_name +'/'+  str(time.time()) + '.jpg'
+                    else:
+                        save_image = None
                     logger.info(f'Processing capture for {self.camera_name}')
-                    model_msgs = self.capture_and_prepare()
+                    model_msgs = self.capture_and_prepare(save_image = save_image)
                     for model_msg in model_msgs:
                         # Use put_nowait for multiprocessing queue to avoid blocking
                         try: