soleil-kamitto/skinscreen

1from datetime import datetime, timezone2 3from sqlalchemy import JSON, Boolean, DateTime, Float, ForeignKey, Integer, String, create_engine4from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, sessionmaker5 6from app.config import settings7 8connect_args = {"check_same_thread": False} if settings.database_url.startswith("sqlite") else {}9engine = create_engine(settings.database_url, connect_args=connect_args)10SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)11 12 13class Base(DeclarativeBase):14    pass15 16 17def utcnow() -> datetime:18    return datetime.now(timezone.utc)19 20 21class User(Base):22    __tablename__ = "users"23 24    id: Mapped[int] = mapped_column(Integer, primary_key=True)25    email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)26    password_hash: Mapped[str] = mapped_column(String(255), nullable=False)27    is_admin: Mapped[bool] = mapped_column(Boolean, default=False)28    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)29 30    patients: Mapped[list["Patient"]] = relationship(back_populates="owner")31 32 33class Patient(Base):34    __tablename__ = "patients"35 36    id: Mapped[int] = mapped_column(Integer, primary_key=True)37    user_id: Mapped[int] = mapped_column(ForeignKey("users.id"), nullable=False)38    name: Mapped[str] = mapped_column(String(255), nullable=False)39    age: Mapped[int] = mapped_column(Integer, nullable=False)40    skin_type: Mapped[str] = mapped_column(String(50), default="unknown")41    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)42 43    owner: Mapped["User"] = relationship(back_populates="patients")44    predictions: Mapped[list["Prediction"]] = relationship(back_populates="patient")45 46 47class Prediction(Base):48    __tablename__ = "predictions"49 50    id: Mapped[int] = mapped_column(Integer, primary_key=True)51    patient_id: Mapped[int] = mapped_column(ForeignKey("patients.id"), nullable=False)52    image_path: Mapped[str] = mapped_column(String(512), nullable=False)53    predicted_class: Mapped[str] = mapped_column(String(50), nullable=False)54    melanoma_prob: Mapped[float] = mapped_column(Float, nullable=False)55    confidence: Mapped[float] = mapped_column(Float, nullable=False)56    probabilities: Mapped[dict] = mapped_column(JSON, nullable=False)57    gradcam_heatmap: Mapped[str] = mapped_column(String, nullable=False)58    requires_referral: Mapped[bool] = mapped_column(Boolean, default=False)59    notes: Mapped[str | None] = mapped_column(String(2000), nullable=True)60    deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)61    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)62 63    patient: Mapped["Patient"] = relationship(back_populates="predictions")64 65 66class SystemSettings(Base):67    __tablename__ = "system_settings"68 69    id: Mapped[int] = mapped_column(Integer, primary_key=True)70    referral_melanoma_threshold: Mapped[float] = mapped_column(Float, default=0.3)71    referral_confidence_threshold: Mapped[float] = mapped_column(Float, default=70.0)72 73 74def get_or_create_settings(db) -> "SystemSettings":75    row = db.query(SystemSettings).filter(SystemSettings.id == 1).first()76    if row is None:77        row = SystemSettings(id=1)78        db.add(row)79        db.commit()80        db.refresh(row)81    return row82 83 84def init_db() -> None:85    Base.metadata.create_all(bind=engine)86 87 88def get_db():89    db = SessionLocal()90    try:91        yield db92    finally:93        db.close()94