views.py

import flask
from flask import Blueprint, Response, render_template, request, jsonify
from camera import Camera
from detectors.color import ColorDetector
from detectors.yolo import YoloDetector
import cv2, threading, numpy as np

bp = Blueprint("main", __name__)
cam = Camera()
lock = threading.Lock()
detectors = {
    "color": ColorDetector(hsv=(60,155,155), tol=(25,100,100), min_area=100, s_floor=10, v_floor=10),
    "yolo": YoloDetector(weights="models/penis.pt", conf=0.24, iou=0.45)
}

mode = {"name": "color"}

@bp.route("/color_debug")
def color_debug():
    f = cam.read()
    if f is None:
        return jsonify(ok=False), 503
    with lock:
        stats = detectors["color"].probe(f)
        hsv = detectors["color"].hsv
        tol = detectors["color"].tol
    return jsonify(ok=True, hsv_center=hsv, tol=tol, space=detectors["color"].space, **stats)

@bp.route("/set_color_space")
def set_color_space():
    sp = request.args.get("space", "both")
    with lock:
        detectors["color"].set_space(sp)
    return jsonify(ok=True, space=sp)

def stream_gen():
    font = cv2.FONT_HERSHEY_SIMPLEX
    while True:
        f = cam.read()
        if f is None:
            break
        with lock:
            m = mode["name"]
            if m == "color":
                out = detectors["color"].process(f.copy())
                c = "C:" + str(len(cv2.findContours(cv2.cvtColor(detectors["color"].mask_frame(f.copy()), cv2.COLOR_BGR2GRAY), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]))
                y = "Y:-"
            elif m == "yolo":
                out = detectors["yolo"].process(f.copy())
                c = "C:-"
                y = "Y:" + str(len(detectors["yolo"].model.predict(f.copy(), conf=detectors["yolo"].conf, iou=detectors["yolo"].iou, imgsz=getattr(detectors["yolo"], "imgsz", 640), verbose=False)[0].boxes))
            else:
                out = detectors["color"].process(f.copy())
                out = detectors["yolo"].process(out)
                c = "C:" + str(len(cv2.findContours(cv2.cvtColor(detectors["color"].mask_frame(f.copy()), cv2.COLOR_BGR2GRAY), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]))
                y = "Y:" + str(len(detectors["yolo"].model.predict(f.copy(), conf=detectors["yolo"].conf, iou=detectors["yolo"].iou, imgsz=getattr(detectors["yolo"], "imgsz", 640), verbose=False)[0].boxes))
        cv2.putText(out, f"MODE: {m.upper()}  {c}  {y}", (12, 28), font, 0.8, (0,255,0), 2, cv2.LINE_AA)
        ok, buf = cv2.imencode(".jpg", out, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
        if not ok:
            continue
        b = buf.tobytes()
        yield (b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: " + f"{len(b)}".encode() + b"\r\n\r\n" + b + b"\r\n")

def mask_gen():
    while True:
        f = cam.read()
        if f is None:
            break
        with lock:
            m = detectors["color"].mask_frame(f.copy())
        ok, buf = cv2.imencode(".jpg", m, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
        if not ok:
            continue
        b = buf.tobytes()
        yield (b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: " + f"{len(b)}".encode() + b"\r\n\r\n" + b + b"\r\n")

@bp.route("/get_mode")
def get_mode():
    with lock:
        m = mode["name"]
    return jsonify(ok=True, mode=m)

@bp.route("/")
def home():
    return render_template("layout.html")

@bp.route("/mjpeg")
def mjpeg():
    return Response(stream_gen(), mimetype="multipart/x-mixed-replace; boundary=frame")

@bp.route("/set_mode")
def set_mode():
    m = request.args.get("m", "color")
    with lock:
        if m in ("color", "yolo", "both"):
            mode["name"] = m
            return jsonify(ok=True, mode=m)
    return jsonify(ok=False, error="invalid mode"), 400

@bp.route("/set_hsv")
def set_hsv():
    h = int(request.args.get("h", 100))
    s = int(request.args.get("s", 200))
    v = int(request.args.get("v", 158))
    dh = int(request.args.get("dh", 25))
    ds = int(request.args.get("ds", 50))
    dv = int(request.args.get("dv", 100))
    with lock:
        c = detectors["color"]
        c.set_hsv(h, s, v, dh, ds, dv)
    return jsonify(ok=True, hsv=[h, s, v], tol=[dh, ds, dv])

@bp.route("/yolo_classes")
def yolo_classes():
    with lock:
        y = detectors["yolo"]
        names = y.class_names()
    return jsonify(ok=True, classes=names)

@bp.route("/set_yolo_labels")
def set_yolo_labels():
    labels = request.args.get("labels", None)
    with lock:
        y = detectors["yolo"]
        y.set_labels(labels)
    return jsonify(ok=True, labels=labels)

@bp.route("/set_yolo_conf")
def set_yolo_conf():
    conf = float(request.args.get("conf", 0.25))
    with lock:
        y = detectors["yolo"]
        y.set_conf(conf)
    return jsonify(ok=True, conf=conf)

@bp.route("/set_yolo_iou")
def set_yolo_iou():
    iou = float(request.args.get("iou", 0.45))
    with lock:
        y = detectors["yolo"]
        y.set_iou(iou)
    return jsonify(ok=True, iou=iou)

@bp.route("/set_yolo_weights")
def set_yolo_weights():
    path = request.args.get("path", "models/best.pt")
    with lock:
        y = detectors["yolo"]
        y.set_weights(path)
    return jsonify(ok=True, weights=path)

@bp.route("/mask")
def mask():
    return Response(mask_gen(), mimetype="multipart/x-mixed-replace; boundary=frame")

@bp.route("/set_min_area")
def set_min_area():
    a = int(request.args.get("a", 5))
    with lock:
        detectors["color"].set_min_area(a)
    return jsonify(ok=True, min_area=a)

@bp.route("/pick")
def pick():
    x = int(request.args.get("x", -1))
    y = int(request.args.get("y", -1))
    box = int(request.args.get("box", 15))
    f = cam.read()
    if f is None or x < 0 or y < 0:
        return jsonify(ok=False), 400
    x0 = max(0, x - box); y0 = max(0, y - box)
    x1 = min(f.shape[1] - 1, x + box); y1 = min(f.shape[0] - 1, y + box)
    roi = f[y0:y1+1, x0:x1+1].copy()
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    H = int(np.median(hsv[..., 0])); S = int(np.median(hsv[..., 1])); V = int(np.median(hsv[..., 2]))
    with lock:
        dh, ds, dv = detectors["color"].autotune_tol(hsv)
        detectors["color"].set_hsv(H, S, V, dh, ds, dv)
    return jsonify(ok=True, hsv=[H, S, V], tol=[dh, ds, dv], box=box, at=[x, y])

@bp.route("/snap_color")
def snap_color():
    f = cam.read()
    if f is None:
        return Response(status=503)
    with lock:
        out = detectors["color"].process(f.copy())
    ok, buf = cv2.imencode(".jpg", out, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
    return Response(buf.tobytes(), mimetype="image/jpeg") if ok else Response(status=500)

@bp.route("/snap_raw")
def snap_raw():
    f = cam.read()
    if f is None:
        return Response(status=503)
    ok, buf = cv2.imencode(".jpg", f, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
    return Response(buf.tobytes(), mimetype="image/jpeg") if ok else Response(status=500)

@bp.route("/set_color_draw")
def set_color_draw():
    b = request.args.get("b"); g = request.args.get("g"); r = request.args.get("r"); th = request.args.get("th")
    with lock:
        detectors["color"].set_draw(
            b=None if b is None else int(b),
            g=None if g is None else int(g),
            r=None if r is None else int(r),
            thickness=None if th is None else int(th)
        )
    return jsonify(ok=True)

@bp.route("/set_color_hex")
def set_color_hex():
    hx = request.args.get("hex", "#00FF00")
    with lock:
        detectors["color"].set_draw_hex(hx)
    return jsonify(ok=True, hex=hx)

@bp.route("/set_yolo_draw")
def set_yolo_draw():
    b = request.args.get("b"); g = request.args.get("g"); r = request.args.get("r"); th = request.args.get("th")
    with lock:
        detectors["yolo"].set_draw(
            b=None if b is None else int(b),
            g=None if g is None else int(g),
            r=None if r is None else int(r),
            thickness=None if th is None else int(th)
        )
    return jsonify(ok=True)

@bp.route("/set_yolo_hex")
def set_yolo_hex():
    hx = request.args.get("hex", "#FF0000")
    with lock:
        detectors["yolo"].set_draw_hex(hx)
    return jsonify(ok=True, hex=hx)