Upload 4 files

app.py

@@ -6,14 +6,13 @@ from ultralytics import YOLO
 from google.cloud import vision
 _api_key = os.environ["API_KEY"]
 _project_id = os.environ["PROJECT_ID"]
-
 client = vision.ImageAnnotatorClient(client_options={"quota_project_id": _project_id, "api_key": _api_key})
 
 import math
 from scipy.spatial import KDTree
 import io
 from time import time
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFilter
 import numpy as np
 import cv2
 
@@ -23,9 +22,10 @@ modelPh = r'corners-best.pt'
 
 model1DIM = 640
 keypointModel = r'keypoints-best.pt'
+minSz = 1280
 
 
-_examples = ["example1.jpg", "example2.jpg"]
+_examples = [["example1.jpg",True], ["example2.jpg",False], ["example3.jpg",True]]
 
 
 def unwarp_image(warped_image, src_points, dst_points, output_width, output_height):
@@ -57,14 +57,18 @@ def get_corners(results:list, img):
     #  keypoints ie corners for homography
     KP = "topLeft topRight bottomRight bottomLeft".split()
     r = results[0]
-    kpco = r.keypoints.xy.cpu().squeeze()
-    assert len(kpco)>0, "not found"
-    keypoints = {k:v.numpy() for v,k in zip(kpco,KP)}
-    sz = model1DIM
-    dstCorners = np.array([(0,0),(sz,0),(sz,sz),(0,sz)])
-    planar = unwarp_image(img, np.array(list(keypoints.values())),dstCorners, sz,sz)
-    # planar.save("temp-ph.jpg")
-    return planar, keypoints
+    planars = []
+    kps = []
+    for kpco in r.keypoints.xy.cpu():#.squeeze()
+        assert len(kpco)>0, "not found"
+        keypoints = {k:v.numpy() for v,k in zip(kpco,KP)}
+        sz = model1DIM
+        dstCorners = np.array([(0,0),(sz,0),(sz,sz),(0,sz)])
+        planar = unwarp_image(img, np.array(list(keypoints.values())),dstCorners, sz,sz)
+        # planar.save("temp-ph.jpg")
+        planars.append(planar)
+        kps.append(keypoints)
+    return planars, kps
 
 
 model = None
@@ -88,6 +92,8 @@ def preprocessImg(planar):
     elif w!=h:
         img = img.resize((_,_))
 
+    if _ < minSz:
+        img = img.resize((minSz,minSz))
     return img
 
 
@@ -174,7 +180,8 @@ def result_as_validvalue(contents:list[dict])->tuple[list[dict], list[str]]:
             continue
         b = f["boxCorners"]
         m = median_point_of_bounding_box(*np.array(b).flatten())
-        valid.append({"text":f["text"], "value": value, "mid": m})
+        a = cv2.contourArea(np.array(b)) / len(f["text"])
+        valid.append({"text":f["text"], "value": value, "mid": m, "apchar":a, "box":b})
 
     valid.sort(key=lambda e: e["value"])
     return valid, other
@@ -200,9 +207,9 @@ def determine_ocr_neighbors(center, valid:list[dict])->tuple[ list, float ]:
         a = np.array(values[-1]["mid"]) - center 
         b = np.array(v["mid"]) - center
         ang,_ = cosangle(a,b)
-        if _ <0:
-            Warning(f"skipping {u['value']} rot:{_}")
-            continue
+        # if _ <0:
+        #     Warning(f"skipping {u['value']} rot:{_}")
+        #     continue
         angS += ang
         u["dang"] = ang
         u["dvda"] = u["dv"] / ang
@@ -250,6 +257,17 @@ def angles_from_tip(keypoints, values, nearestIx):
     return values
 
 
+def sort_clockwise_with_start(coordinates, x_center, y_center, starting_index):
+    angles = [math.atan2(y - y_center, x - x_center) for x, y in coordinates]
+    sorted_indices = sorted(range(len(angles)), key=lambda i: (angles[i] - angles[starting_index] + 2 * math.pi) % (2 * math.pi))
+    return sorted_indices, angles
+
+def remove_nonrange_value(valid):
+    meanArea = np.mean([e["apchar"] for e in valid])
+    cutoff = 0.5
+    valid = list(filter(lambda e: abs(e["apchar"]-meanArea)/meanArea < cutoff, valid))
+    return valid
+
 def get_needle_value(img, keypoints):
 
     tic2 = time()
@@ -260,18 +278,32 @@ def get_needle_value(img, keypoints):
     assert len(contents)
     valid,other = result_as_validvalue(contents)
     assert len(valid)
+
+    valid.append({"text":"tip", "mid":keypoints["tip"]})
+    ix,an = sort_clockwise_with_start([e["mid"] for e in valid],*keypoints["center"], 0)
+    valid = [valid[i] for i in ix]
+    assert valid[-1]["text"]!="tip" and valid[0]["text"]!="tip", "failed to properly detect tip"
+    nearestIx=[]
+    for i,v in enumerate(valid):
+        if "tip"==v["text"]:
+            nearestIx = [i-1,i]
+            valid.pop(i)
+            break
+    nearestIx = np.array(nearestIx)
+    valid = remove_nonrange_value(valid)
+
     center = np.array(keypoints["center"])
     values, rate = determine_ocr_neighbors(center, valid)
     assert len(values)>=2, "failed to find at least 2 OCR values"
 
     # import pandas as pd
     # print(pd.DataFrame.from_dict(values))
-    
-    tree = KDTree([v["mid"] for v in values])
-    # find bounding ocr values of tip
-    dist, nearestIx = tree.query(keypoints["tip"],k=2)
-    nearestIx.sort()
-    dist, nearestIx 
+
+    # tree = KDTree([v["mid"] for v in values])
+    # # find bounding ocr values of tip
+    # dist, nearestIx = tree.query(keypoints["tip"],k=2)
+    # nearestIx.sort()
+    # dist, nearestIx 
 
     values = angles_from_tip(keypoints, values, nearestIx)
     # compare against start and end
@@ -336,36 +368,61 @@ def predict(img, detect_gauge_first):
     if detect_gauge_first:
         model0 = get_load_PhModel()
         results = model0.predict(img)
-        phimg,_ = get_corners(results, img)
+        phimgs,_ = get_corners(results, img)
+        if len(phimgs)==0:
+            raise gr.Error("no gauge found")
     else:
-        phimg = img.copy()
+        phimgs = [img.copy()]
 
-    model = get_load_KpModel()
-    phimg = preprocessImg(phimg)
-    results = model.predict(phimg)
-    keypoints = get_keypoints(results)
+    payloads = []
+    for phimg in phimgs:
+        model = get_load_KpModel()
+        phimg = preprocessImg(phimg)
+        results = model.predict(phimg)
+        keypoints = get_keypoints(results)
 
-    angle2tip, totalAngle = calculate_sweep_angles(keypoints)
+        angle2tip, totalAngle = calculate_sweep_angles(keypoints)
 
-    payload = get_needle_value(phimg, keypoints)
-    payload["angleToTip"] = round(float(angle2tip),2)
-    payload["totalAngle"] = round(float(totalAngle),2)
-    for k,v in payload.items():
-        print(k, type(v),v)
+        phimg = phimg.filter(ImageFilter.UnsharpMask(radius=3))
+        payload = get_needle_value(phimg, keypoints)
+        payload["angleToTip"] = round(float(angle2tip),2)
+        payload["totalAngle"] = round(float(totalAngle),2)
+        for k,v in payload.items():
+            print(k, type(v),v)
+        payloads.append(payload)
 
-    return payload
+    return payloads
 
 
 def test(img, detect_gauge_first):
     return {"msg":str(img.size), "other": detect_gauge_first}
 
-
-gr.Interface(fn=predict,
+description = r"""
+<b>Official 🤗 Gradio demo</b> for <a href='https://synanthropic.com/reading-analog-gauge' target='_blank'><b>Reading Analog Gauges: Automate Gauge Readings with AI in Days, Not Months
+</b></a>.<br>
+<br>
+This model reads analog dial gauge by detecting, applying perspective correction, and gauge reading. 
+<br>
+The model was build <i><strong>only</strong></i> with synthetic data.<br>
+<br>
+You can read more about it [here](https://synanthropic.com/reading-analog-gauge).
+<br>
+<br>
+❗️Usage steps:<br>
+1️⃣ Upload an image with analog dial gauge with readable values. The gauge face in the uploaded image should <b>occupy the majority of the image</b>.<br>
+2️⃣ If the image has only one gauge and is a direct flat view, uncheck <strong>detect gauge first</strong>.</br>
+3️⃣ Click the <b>Submit</b> button to start inference.<br>
+<br>
+"""
+
+gr.Interface(title="title",
+             description=description,
+             fn=predict,
              inputs=[
-                 gr.Image(type="pil", sources=["upload","clipboard"],streaming=False, min_width=640),
+                 gr.Image(type="pil", sources=["upload"],streaming=False, min_width=640),
                  gr.Checkbox(True, label="detect gauge first", info="if input image is zoomed in on only one gauge, uncheck box")
                  ],
              outputs="json",
-             examples=[_examples],
+             examples=_examples,
              cache_examples=True)\
     .launch()