Refactor storage

app.py

@@ -1,110 +1,96 @@
-import os
 import json
 from PIL import Image
 from src.digitizer import Digitizer
 from src.routing import IdentifierConverter
 from src.draw import draw_objects
+from src.file_storage import FileStorage
 from io import BytesIO
 from glob import glob
-from flask import Flask, request
+from flask import Flask, request, send_file
 
 model = 'models/yolov8-detect-20240229.onnx'
 data_dir = 'data/'
 
 digitizer = Digitizer(model)
+storage = FileStorage(data_dir)
 
 app = Flask(__name__)
 app.url_map.converters['identifier'] = IdentifierConverter
 
 @app.route('/digitize', methods=['POST'])
 def digitize():
-    value, _, _ = _detect()
+    decimals = request.args.get('decimals', default=0, type=int)
+    threshold = request.args.get('threshold', default=0.7, type=float)
+
+    image = Image.open(BytesIO(request.get_data()))
+    value, _ = digitizer.detect(image, decimals, threshold)
 
     if value is None:
         return 'No reading found', 400
 
-    json_data = json.dumps({'value': value})
-
-    return json_data, 200, {'Content-Type': 'application/json'}
+    return json.dumps({'value': value}), 200, {'Content-Type': 'application/json'}
 
 @app.route('/meter/<identifier:meter_id>', methods=['POST'])
 def update_meter(meter_id):
-    value, objects, image = _detect()
+    decimals = request.args.get('decimals', default=0, type=int)
+    threshold = request.args.get('threshold', default=0.7, type=float)
+    max_increase = request.args.get('max_increase', default=float('inf'), type=float)
+
+    image = Image.open(BytesIO(request.get_data()))
+    value, objects = digitizer.detect(image, decimals, threshold)
 
     if value is None:
         return 'No reading found', 400
 
-    json_file = os.path.join(data_dir, f'{meter_id}.json')
-    max_increase = request.args.get('max_increase', default=float('inf'), type=float)
-
-    if os.path.exists(json_file):
-        with open(json_file, 'r') as f:
-            data = json.load(f)
-        if value < data['value']:
+    try:
+        old_value = storage.read_value(meter_id)
+
+        if value < old_value:
             return 'Reading is lower than previous', 400
-        if value - data['value'] > max_increase:
+        if value - old_value > max_increase:
             return 'Reading increased too much', 400
+    except FileNotFoundError:
+        # no old value
+        pass
 
-    json_data = json.dumps({'value': value})
-
-    with open(json_file, 'w') as f:
-        f.write(json_data)
-
-    image.save(os.path.join(data_dir, f'{meter_id}.jpg'))
-
-    draw_objects(image, objects)
-
-    image.save(os.path.join(data_dir, f'{meter_id}_result.jpg'))
+    result_image = image.copy() 
+    draw_objects(result_image, objects)
+
+    storage.store(meter_id, value, image, result_image)
 
-    return json_data, 200, {'Content-Type': 'application/json'}
+    return json.dumps({'value': value}), 200, {'Content-Type': 'application/json'}
 
 @app.route('/meter/<identifier:meter_id>', methods=['GET'])
 def show_meter(meter_id):
-    return _send_file(f'{meter_id}.json', 'application/json')
+    try:
+        value = storage.read_value(meter_id)
+        return json.dumps({'value': value}), 200, {'Content-Type': 'application/json'}
+    except FileNotFoundError:
+        return 'Meter not found', 404
 
 @app.route('/meter/<identifier:meter_id>/image', methods=['GET'])
 def meter_image(meter_id):
-    return _send_file(f'{meter_id}.jpg', 'image/jpeg')
+    try:
+        data = storage.read_image_data(meter_id)
+        return data, 200, {'Content-Type': 'image/jpeg'} 
+    except FileNotFoundError:
+        return 'Meter not found', 404
 
-@app.route('/meter/<identifier:meter_id>/image_result', methods=['GET'])
-def meter_image_result(meter_id):
-    return _send_file(f'{meter_id}_result.jpg', 'image/jpeg')
+@app.route('/meter/<identifier:meter_id>/result_image', methods=['GET'])
+def meter_result_image(meter_id):
+    try:
+        data = storage.read_result_image_data(meter_id)
+        return data, 200, {'Content-Type': 'image/jpeg'} 
+    except FileNotFoundError:
+        return 'Meter not found', 404
 
 @app.route('/meter/<identifier:meter_id>/reset', methods=['GET'])
 def reset_meter(meter_id):
-    for file in glob(f'{data_dir}{meter_id}.*'):
-        os.remove(file)
-
     value = request.args.get('value', default=None, type=float)
+    storage.remove(meter_id)
 
     if value is not None:
-        json_data = json.dumps({'value': value})
-        json_file = os.path.join(data_dir, f'{meter_id}.json')
-        with open(json_file, 'w') as f:
-            f.write(json_data)
+        blank_image = Image.new('RGB', (1, 1))
+        storage.store(meter_id, value, blank_image, blank_image)
 
     return 'Meter reset', 200
-
-def _detect():
-    decimals = request.args.get('decimals', default=0, type=int)
-    threshold = request.args.get('threshold', default=0.7, type=float)
-    data = request.get_data()
-    image = Image.open(BytesIO(data))
-    reading, objects = digitizer.detect(image, threshold)
-
-    if len(reading):
-        value = float(reading) / (10 ** decimals)
-    else:
-        value = None
-
-    return value, objects, image
-
-def _send_file(file_name, content_type):
-    path = os.path.join(data_dir, file_name)
-
-    if not os.path.exists(path):
-        return 'Meter not found', 404
-
-    data = open(path, 'rb').read()
-
-    return data, 200, {'Content-Type': content_type}

src/digitizer.py

@@ -8,7 +8,7 @@ class Digitizer:
     def __init__(self, model_file):
         self.model = YOLOv8(model_file)
 
-    def detect(self, image, conf_threshold=DEFAULT_THRESHOLD):
+    def detect(self, image, decimals=0, conf_threshold=DEFAULT_THRESHOLD):
         output = self.model.detect_objects(image, conf_threshold)
 
         objects = [DetectedObject(*r) for r in zip(*output)]
@@ -18,9 +18,9 @@ def detect(self, image, conf_threshold=DEFAULT_THRESHOLD):
         objects = sorted(objects, key=lambda o: o.box[0])
         objects = self._remove_overlapping(objects)
 
-        reading = ''.join([str(o.class_id) for o in objects])
+        value = self._digitize(objects, decimals)
 
-        return reading, objects
+        return value, objects
 
     # remove overlapping boxes by x coordinate keeping the one with the highest score
     def _remove_overlapping(self, objects):
@@ -35,4 +35,12 @@ def _remove_overlapping(self, objects):
             elif objects[i].score > output[-1].score:
                 output[-1] = objects[i]
 
-        return output
+        return output
+
+    def _digitize(self, objects, decimals):
+        string = ''.join([str(o.class_id) for o in objects])
+
+        if len(string) > 0:
+            return float(string) / (10 ** decimals)
+        else:
+            return None

src/file_storage.py

@@ -0,0 +1,42 @@
+import json
+import os
+
+class FileStorage:
+    def __init__(self, path):
+        self.path = path
+
+    def store(self, meter_id, value, image, result_image):
+        self._write_data(f'{meter_id}.json', {'value': value})
+        self._write_image(f'{meter_id}.jpg', image)
+        self._write_image(f'{meter_id}_result.jpg', result_image)
+
+    def read_value(self, meter_id):
+        return self._read_data(f'{meter_id}.json')['value']
+
+    def read_image_data(self, meter_id):
+        return self._read_image_data(f'{meter_id}.jpg')
+
+    def read_result_image_data(self, meter_id):
+        return self._read_image_data(f'{meter_id}_result.jpg')
+
+    def remove(self, meter_id):
+        files = [f'{meter_id}.json', f'{meter_id}.jpg', f'{meter_id}_result.jpg']
+        for file in files:
+            path = os.path.join(self.path, file)
+            if os.path.exists(path):
+                os.remove(path)
+
+    def _write_data(self, name, data):
+        with open(os.path.join(self.path, name), 'w') as f:
+            json.dump(data, f)
+
+    def _read_data(self, name):
+        with open(os.path.join(self.path, name), 'r') as f:
+            return json.load(f)
+
+    def _write_image(self, name, image):
+        with open(os.path.join(self.path, name), 'wb') as f:
+            image.save(f, 'JPEG')
+
+    def _read_image_data(self, name):
+        return open(os.path.join(self.path, name), 'rb').read()