Update app.py

app.py

@@ -1,8 +1,8 @@
 import base64
 import cv2
 import numpy as np
-import os
-from fastapi import FastAPI
+import uvicorn
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from collections import defaultdict
 
@@ -19,25 +19,30 @@ def root():
 class Input(BaseModel):
     image_base64: str
 
-def save_base64_image_cv(base64_str, output_path="final.png"):
-    img_data = base64.b64decode(base64_str)
-    nparr = np.frombuffer(img_data, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
-
-    if img.shape[2] == 4:
-        alpha = img[:, :, 3] / 255.0
-        rgb = img[:, :, :3]
-        white_bg = np.ones_like(rgb, dtype=np.uint8) * 255
-        img = (rgb * alpha[:, :, None] + white_bg * (1 - alpha[:, :, None])).astype(np.uint8)
+def preprocess_image_memory(base64_str):
+    try:
+        img_data = base64.b64decode(base64_str)
+        nparr = np.frombuffer(img_data, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
+        
+        if img is None:
+            raise ValueError("Invalid/corrupt image")
 
-    cv2.imwrite(output_path, img)
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            alpha = img[:, :, 3] / 255.0
+            rgb = img[:, :, :3]
+            white_bg = np.ones_like(rgb, dtype=np.uint8) * 255
+            img = (rgb * alpha[:, :, None] + white_bg * (1 - alpha[:, :, None])).astype(np.uint8)
+        
+        return img
+    except Exception as e:
+        raise ValueError(f"Error decoding image: {str(e)}")
 
-def extract_icon_positions(image_path):
-    img = cv2.imread(image_path)
+def extract_icon_positions(img):
     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     _, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
     contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
+    
     icons, pos = [], []
     for c in contours:
         x, y, w, h = cv2.boundingRect(c)
@@ -52,16 +57,24 @@ def img_hash(img):
     return (img > img.mean()).astype(np.uint8).flatten()
 
 def find_rarest(icon_features, positions):
+    if not icon_features:
+        return None, None
+        
     hashes = [img_hash(i) for i in icon_features]
     groups = defaultdict(list)
-
+    
     for i, h in enumerate(hashes):
-        for g in groups.values():
-            if np.sum(h != hashes[g[0]]) < 3:
-                g.append(i)
+        found = False
+        for label, group in groups.items():
+            if np.sum(h != hashes[group[0]]) < 3:
+                group.append(i)
+                found = True
                 break
-        else:
+        if not found:
             groups[len(groups)] = [i]
+            
+    if not groups:
+        return None, None
 
     idx = min(groups.values(), key=len)[0]
     return positions[idx]
@@ -69,10 +82,19 @@ def find_rarest(icon_features, positions):
 @app.post("/solve")
 def solve(data: Input):
     try:
-        save_base64_image_cv(data.image_base64)
-        icons, pos = extract_icon_positions("final.png")
+        img = preprocess_image_memory(data.image_base64)
+
+        icons, pos = extract_icon_positions(img)
+        
+        if not icons:
+            return {"error": "No icons found", "x": 0, "y": 0}
+
         x, y = find_rarest(icons, pos)
-        return {"x": x, "y": y}
-    finally:
-        if os.path.exists("final.png"):
-            os.remove("final.png")
+
+        return {"x": int(x), "y": int(y)}
+        
+    except Exception as e:
+        return {"error": str(e)}
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)