Files
2026-06-21 21:24:47 +02:00

164 lines
5.5 KiB
Python

from tinydb import Query, TinyDB
from logging import getLogger
from base64 import b64encode
from requests import post, get, RequestException
import subprocess
import time
from os import listdir, path
from re import search
logger = getLogger("hairloss")
OLLAMA_URL = "http://localhost:11434"
PROMPT = """
Think
You are a visual estimator of scalp hair density.
You will be given:
an image
an angle label: "top", "lateral", or "face"
Angle meaning:
top: evaluate crown/vertex only
lateral: evaluate temple recession and side density
face: evaluate frontal hairline and symmetry
Task:
Estimate hair loss severity as a continuous value between 0.0 and 1.0.
Scoring meaning:
0.0 → full dense hair
0.5 → moderate thinning / visible scalp
1.0 → severe hair loss
Guidelines:
Use only visible evidence in the image.
Be robust to lighting and hairstyle, but consider scalp visibility.
Only evaluate regions that are visible from the given angle.
If a region is not visible, do not infer it.
If image quality is unclear:
use a neutral estimate based on visible areas (do not guess extremes)
Output:
Return ONLY a single float between 0.0 and 1.0.
No text, no explanation, no punctuation.
"""
def prepareModel(timeout: int = 15):
try:
get(f"{OLLAMA_URL}/api/tags", timeout=2)
logger.info("Ollama already running")
return
except RequestException:
logger.info("Ollama not running, starting server...")
subprocess.Popen(
["ollama", "serve"],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
start = time.time()
while time.time() - start < timeout:
try:
get(f"{OLLAMA_URL}/api/tags", timeout=2)
logger.info("Ollama server started")
return
except RequestException:
time.sleep(1)
raise RuntimeError("Failed to start Ollama server")
def evaluate(image_path: str, model: str, repetitions: int = 10, timeOut: int = 10):
db = TinyDB(path.abspath(path.join(image_path, "../hairlineResults.json")))
q = Query()
prepareModel()
for image in listdir(image_path):
oldInfo = db.get(q.filename == image)
if model in oldInfo.keys():
logger.info(f"Skipping evaluation {image} with {model}...")
continue
logger.info(f"Evaluating {image} with model {model} {repetitions} times.")
with open(path.join(image_path, image), "rb") as f:
imageb64 = b64encode(f.read()).decode("utf-8")
average = 0
averageWithoutUnsure = 0
averageWithoutUnsureCount = 0
if "f" in image:
angle = "face"
elif "l" in image:
angle = "lateral"
elif "t" in image:
angle = "top"
else:
raise ValueError(f"no angle information in filename {image}")
imagePrompt = PROMPT + f"\nlabel:\"{angle}\""
for i in range(repetitions): # we ask multiple times and get the average response
response = post(
f"{OLLAMA_URL}/api/generate",
json={
"model": model,
"prompt": imagePrompt,
"images": [imageb64],
"stream": False,
"thinking": model != "qwen3-vl:8b" # for some reason this model thinkgs so much, that it spends all it's tokens on thinking and None in output...
},
timeout=timeOut,
)
response.raise_for_status()
logger.debug(response.json())
result = search(r"^(0(?:\.\d+)?|1(?:\.0+)?)$", response.json()["response"].strip())
logger.debug(result)
if result == None:
result = 0.5
else:
data = float(search(r"^(0(?:\.\d+)?|1(?:\.0+)?)$", response.json()["response"].strip()).group(1)) # extract the float in case it outputed some text
if data != 0.5: # if the model is uncertain it should return 0.5
averageWithoutUnsure += data
averageWithoutUnsureCount += 1
average += data
average /= repetitions
if averageWithoutUnsureCount != 0: # avoids case where all repetitions are unsure
averageWithoutUnsure /= averageWithoutUnsureCount
else:
averageWithoutUnsure = 0.5
if db.contains(q.filename == image): # we add the result to a database. It allows to do everything in multiple runs and combine results from multiple models
imageResult = db.get(q.filename == image)
imageResult[model] = average
imageResult[model+"WithoutUnsure"] = averageWithoutUnsure
db.upsert(imageResult, q.filename == image)
else:
imageResult = {
"filename": image,
model: average,
model + "WithoutUnsure": averageWithoutUnsure
}
db.upsert(imageResult, q.filename == image)
logger.info("Normalizing values ...")
# normalise values
maxValue = max(image[model + "WithoutUnsure"] for image in db)
minValue = min(image[model + "WithoutUnsure"] for image in db)
for image in listdir(image_path):
unnormalizedValues = db.get(q.filename == image)
unnormalizedValues[model+"Normalized"] = (unnormalizedValues[model+"WithoutUnsure"] - minValue)/(maxValue - minValue)
db.upsert(unnormalizedValues, q.filename == image)
logger.info("Stopping model...")
subprocess.Popen(
["ollama", "stop", model],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)