ml_models/src/main.py - master - sofnuk2004/web-server

ml_models/src/main.py
122 строки4 KB
sofya365
начало
16 сен 2025, 16:55
16 сен 2025, 16:55c29a2d1
from flask import Flask, request, jsonify
from flask_cors import CORS
import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
import os
import uuid

app = Flask(__name__)
CORS(app)

# Создаем директории для хранения файлов
os.makedirs('uploads', exist_ok=True)
os.makedirs('processed', exist_ok=True)

# Простая заглушка для преобразования RGB -> pseudo-HSI
def rgb_to_pseudo_hsi(image_path):
    """Преобразование RGB в псевдогиперспектральное изображение"""
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Простое преобразование для демонстрации
    hsi = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
    
    # Сохраняем результат
    output_path = f'processed/{uuid.uuid4()}_hsi.png'
    cv2.imwrite(output_path, hsi)
    
    return output_path

# Заглушка для модели ViT
class SimpleViTModel:
    def __init__(self):
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    
    def predict(self, image_path):
        """Простая заглушка для предсказания"""
        image = Image.open(image_path).convert('RGB')
        image = self.transform(image).unsqueeze(0)
        
        # Имитация предсказания
        probabilities = torch.softmax(torch.randn(1, 7), dim=1)
        confidence, diagnosis_idx = torch.max(probabilities, 1)
        
        diagnoses = [
            'melanoma', 'basal cell carcinoma', 'squamous cell carcinoma',
            'nevus', 'seborrheic keratosis', 'dermatofibroma', 'vascular lesion'
        ]
        
        return {
            'diagnosis': diagnoses[diagnosis_idx.item()],
            'confidence': confidence.item(),
            'probabilities': probabilities.tolist()[0]
        }

# Инициализация модели
model = SimpleViTModel()

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy'})

@app.route('/convert_to_hsi', methods=['POST'])
def convert_to_hsi():
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400
    
    # Сохраняем файл
    filename = f'uploads/{uuid.uuid4()}_{file.filename}'
    file.save(filename)
    
    # Конвертируем в HSI
    hsi_path = rgb_to_pseudo_hsi(filename)
    
    return jsonify({
        'hsi_image_path': hsi_path,
        'message': 'Conversion successful'
    })

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400
    
    file = request.files['file']
    model_name = request.form.get('model_name', 'ViT')
    
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400
    
    # Сохраняем файл
    filename = f'uploads/{uuid.uuid4()}_{file.filename}'
    file.save(filename)
    
    # Делаем предсказание
    try:
        result = model.predict(filename)
        
        # Конвертируем в HSI для визуализации
        hsi_path = rgb_to_pseudo_hsi(filename)
        
        return jsonify({
            **result,
            'hsi_image_path': hsi_path,
            'model_used': model_name
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=True)