mirror of
https://github.com/tomasriveral/Hairloss-report.git
synced 2026-08-11 18:38:37 +02:00
164 lines
5.5 KiB
Python
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,
|
|
)
|