init comit

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/vcs.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings" defaultProject="true" />
+</project>

Dockerfile

@@ -0,0 +1,18 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+# Создаем папку для модели
+RUN mkdir -p /home/user/.cache/models
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,103 @@
+import os
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.openapi.docs import get_swagger_ui_html
+from typing import List, Optional
+import io
+from PIL import Image
+import numpy as np
+
+from model import SkinClassifier
+
+# Инициализация FastAPI
+app = FastAPI(
+    title="Skin Classification API",
+    description="API для классификации кожных заболеваний (clear, acne, ros, black)",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
+
+# Инициализация модели
+try:
+    # Путь к модели в Hugging Face Spaces
+    MODEL_PATH = "model/stage1_skin_classifier.pth"
+    classifier = SkinClassifier(MODEL_PATH)
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    classifier = None
+
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    """Главная страница с информацией об API"""
+    return "go to /docs"
+
+
+@app.post("/predict")
+async def predict(
+        file: UploadFile = File(..., description="Изображение для классификации (JPG, PNG)"),
+        return_image: Optional[bool] = False
+):
+    """
+    Классификация одного изображения
+
+    - **file**: Изображение в формате JPG или PNG
+    - **return_image**: Возвращать ли base64 изображение (по умолчанию False)
+    """
+    if classifier is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    # Проверка формата файла
+    allowed_extensions = {".jpg", ".jpeg", ".png", ".bmp"}
+    file_ext = os.path.splitext(file.filename)[1].lower()
+
+    if file_ext not in allowed_extensions:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported file format. Allowed: {allowed_extensions}"
+        )
+
+    try:
+        # Чтение файла
+        contents = await file.read()
+
+        # Предсказание
+        result = classifier.predict(contents)
+
+        response_data = {
+            "filename": file.filename,
+            "prediction": result["predicted_class"],
+            "confidence": result["confidence"],
+            "probabilities": result["all_probabilities"]
+        }
+
+        # Добавляем base64 изображение если нужно
+        if return_image:
+            from io import BytesIO
+            import base64
+
+            img = Image.open(BytesIO(contents))
+            buffered = BytesIO()
+            img.save(buffered, format="JPEG" if file_ext in {".jpg", ".jpeg"} else "PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            response_data["image_base64"] = f"data:image/jpeg;base64,{img_str}"
+
+        return response_data
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
+
+
+# Middleware для обработки CORS
+from fastapi.middleware.cors import CORSMiddleware
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)

model.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import timm
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import numpy as np
+from PIL import Image
+import io
+
+CLASSES = ["clear", "acne", "ros", "black"]
+IMG_SIZE = 224
+
+
+class SkinClassifier:
+    def __init__(self, model_path="model/stage1_skin_classifier.pth"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.classes = CLASSES
+        self.img_size = IMG_SIZE
+
+        # Инициализируем модель
+        self.model = timm.create_model(
+            "efficientnet_b0",
+            pretrained=False,
+            num_classes=len(self.classes)
+        )
+
+        # Загружаем веса
+        state_dict = torch.load(model_path, map_location=self.device)
+        self.model.load_state_dict(state_dict)
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Трансформации
+        self.transform = A.Compose([
+            A.Resize(self.img_size, self.img_size),
+            A.Normalize(
+                mean=(0.485, 0.456, 0.406),
+                std=(0.229, 0.224, 0.225)
+            ),
+            ToTensorV2()
+        ])
+
+    def preprocess(self, image):
+        """Препроцессинг изображения"""
+        if isinstance(image, bytes):
+            image = Image.open(io.BytesIO(image)).convert("RGB")
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image).convert("RGB")
+        else:
+            image = image.convert("RGB")
+
+        image = np.array(image)
+        transformed = self.transform(image=image)
+        return transformed["image"]
+
+    def predict(self, image):
+        """Предсказание класса"""
+        # Препроцессинг
+        tensor = self.preprocess(image)
+        tensor = tensor.unsqueeze(0).to(self.device)
+
+        # Предсказание
+        with torch.no_grad():
+            outputs = self.model(tensor)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+            prediction = torch.argmax(probabilities, dim=1)
+
+            # Получаем вероятности для всех классов
+            probs = probabilities[0].cpu().numpy()
+            class_probs = {self.classes[i]: float(probs[i]) for i in range(len(self.classes))}
+
+        return {
+            "predicted_class": self.classes[prediction.item()],
+            "confidence": float(probabilities[0][prediction.item()]),
+            "all_probabilities": class_probs
+        }

model/stage1_skin_classifier.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:736db77bb261eb8502d54e90a45d20d68b79d6c7852c689cc141fb9411030c32
+size 16351009

requirements.txt

@@ -0,0 +1,14 @@
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+torch==2.1.0
+torchvision==0.16.0
+timm==0.9.12
+scikit-learn==1.3.2
+albumentations==1.3.1
+Pillow==10.1.0
+numpy==1.24.3
+pandas==2.1.4
+tqdm==4.66.1
+opencv-python-headless==4.8.1.78
+requests==2.31.0
+python-multipart==0.0.6