ASU3 / app.py
XORE21's picture
Update app.py
e79654b verified
raw
history blame
2.22 kB
import base64
import cv2
import numpy as np
import os
from fastapi import FastAPI
from pydantic import BaseModel
from collections import defaultdict
app = FastAPI()
@app.get("/")
def root():
return {
"status": "ok",
"service": "iconCaptcha solver",
"endpoint": "/solve"
}
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)
cv2.imwrite(output_path, img)
def extract_icon_positions(image_path):
img = cv2.imread(image_path)
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)
if w > 10 and h > 10:
roi = cv2.resize(thresh[y:y+h, x:x+w], (50, 50))
icons.append(roi)
pos.append((x, y))
return icons, pos
def img_hash(img):
img = cv2.resize(img, (8, 8))
return (img > img.mean()).astype(np.uint8).flatten()
def find_rarest(icon_features, positions):
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)
break
else:
groups[len(groups)] = [i]
idx = min(groups.values(), key=len)[0]
return positions[idx]
@app.post("/solve")
def solve(data: Input):
try:
save_base64_image_cv(data.image_base64)
icons, pos = extract_icon_positions("final.png")
x, y = find_rarest(icons, pos)
return {"x": x, "y": y}
finally:
if os.path.exists("final.png"):
os.remove("final.png")