multi_dataset_api.py

from typing import List, Optional
from fastapi import FastAPI, Query, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from enum import Enum
from datetime import datetime

app = FastAPI(title="Cancer Dataset API")

# Enums for filter options
class DataCategory(str, Enum):
    SKIN = "skin"

class Dataset(str, Enum):
    HAM10000 = "HAM10000"
    ISIC2024 = "ISIC2024"

class Sex(str, Enum):
    MALE = "male"
    FEMALE = "female"

# Models
class Entry(BaseModel):
    id: str
    dataset: str
    disease_type: str
    image_path: str
    sex: Sex
    age: int
    location: str

class FilterType(str, Enum):
    SINGLE_SELECT = "single_select"
    MULTI_SELECT = "multi_select"
    NUMERIC_RANGE = "numeric_range"

class SelectOption(BaseModel):
    value: str
    label: str

class SelectFilterOption(BaseModel):
    type: FilterType
    options: List[SelectOption]

class NumericRangeFilterOption(BaseModel):
    type: FilterType
    min_value: int
    max_value: int

FilterOption = SelectFilterOption | NumericRangeFilterOption

class FilterConfig(BaseModel):
    data_category: FilterOption = Field(description="Data category filter")
    disease_types: FilterOption = Field(description="Disease types filter")
    datasets: FilterOption = Field(description="Datasets filter")
    locations: FilterOption = Field(description="Locations filter")
    sex: FilterOption = Field(description="Sex filter")
    age: FilterOption = Field(description="Age filter")

class PaginatedResponse(BaseModel):
    total: int
    page: int
    page_size: int
    items: List[Entry]

# Example data
EXAMPLE_ENTRIES = [
    Entry(
        id="1",
        dataset="HAM10000",
        disease_type="melanoma",
        image_path="/data/images/001.jpg",
        sex=Sex.MALE,
        age=45,
        location="back"
    ),
    Entry(
        id="2",
        dataset="ISIC2024",
        disease_type="basal cell carcinoma",
        image_path="/data/images/002.jpg",
        sex=Sex.FEMALE,
        age=32,
        location="face"
    )
]

# Filter configuration endpoint
@app.get("/entry/filter-config", response_model=FilterConfig)
async def get_filter_config():
    return FilterConfig(
        data_category=SelectFilterOption(
            type=FilterType.SINGLE_SELECT,
            options=[
                SelectOption(value=cat.value, label=cat.value.title())
                for cat in DataCategory
            ]
        ),
        disease_types=SelectFilterOption(
            type=FilterType.MULTI_SELECT,
            options=[
                SelectOption(value="melanoma", label="Melanoma"),
                SelectOption(value="basal cell carcinoma", label="Basal Cell Carcinoma"),
                SelectOption(value="squamous cell carcinoma", label="Squamous Cell Carcinoma")
            ]
        ),
        datasets=SelectFilterOption(
            type=FilterType.MULTI_SELECT,
            options=[
                SelectOption(value=ds.value, label=ds.value)
                for ds in Dataset
            ]
        ),
        locations=SelectFilterOption(
            type=FilterType.MULTI_SELECT,
            options=[
                SelectOption(value="head_neck", label="Head and Neck"),
                SelectOption(value="upper_extremity", label="Upper Extremity"),
                SelectOption(value="torso", label="Torso")
            ]
        ),
        sex=SelectFilterOption(
            type=FilterType.MULTI_SELECT,
            options=[
                SelectOption(value=sex.value, label=sex.value.title())
                for sex in Sex
            ]
        ),
        age=NumericRangeFilterOption(
            type=FilterType.NUMERIC_RANGE,
            min_value=0,
            max_value=100
        )
    ).model_dump(exclude_none=True, exclude_unset=True)

# List entries endpoint with filtering
@app.get("/entry", response_model=PaginatedResponse)
async def list_entries(
    page: int = Query(1, ge=1),
    page_size: int = Query(10, ge=1, le=100),
    data_category: Optional[str] = None,
    disease_types: List[str] = Query(None),
    sex: List[Sex] = Query(None),
    locations: List[str] = Query(None),
    age_min: Optional[int] = Query(None, ge=0),
    age_max: Optional[int] = Query(None, le=100),
    datasets: List[str] = Query(None)
):
    # Filter logic (simplified for example)
    filtered_entries = EXAMPLE_ENTRIES.copy()
    
    if data_category:
        filtered_entries = [e for e in filtered_entries if data_category.lower() in e.dataset.lower()]
    
    if disease_types:
        filtered_entries = [e for e in filtered_entries if e.disease_type in disease_types]
        
    if sex:
        filtered_entries = [e for e in filtered_entries if e.sex in sex]
        
    if locations:
        filtered_entries = [e for e in filtered_entries if e.location in locations]
        
    if age_min is not None:
        filtered_entries = [e for e in filtered_entries if e.age >= age_min]
        
    if age_max is not None:
        filtered_entries = [e for e in filtered_entries if e.age <= age_max]
        
    if datasets:
        filtered_entries = [e for e in filtered_entries if e.dataset in datasets]

    # Pagination
    total = len(filtered_entries)
    start_idx = (page - 1) * page_size
    end_idx = start_idx + page_size
    page_entries = filtered_entries[start_idx:end_idx]

    return PaginatedResponse(
        total=total,
        page=page,
        page_size=page_size,
        items=page_entries
    )

# Download endpoint
@app.post("/entry/download")
async def download_entries(
    data_category: Optional[str] = None,
    disease_types: List[str] = Query(None),
    sex: List[Sex] = Query(None),
    locations: List[str] = Query(None),
    age_min: Optional[int] = Query(None, ge=0),
    age_max: Optional[int] = Query(None, le=100),
    datasets: List[str] = Query(None)
):
    
    content = "Example export file with filtered data\n"
    content += f"Generated at: {datetime.now().isoformat()}\n"
    content += f"Filters applied: {locals()}"
    
    with open("/tmp/export.txt", "w") as f:
        f.write(content)
    
    return FileResponse(
        "/tmp/export.txt",
        media_type="text/plain",
        filename="cancer_dataset_export.txt"
    )