from fastapi import FastAPI, File, UploadFile, Response
from fastapi.responses import FileResponse
from tensorflow.keras.preprocessing.image import img_to_array
import tensorflow as tf
import cv2
import numpy as np
import os
from scipy import fftpack
from scipy import ndimage
from ultralytics import YOLO
from PIL import Image
import io

app = FastAPI()
uploads_dir = 'uploads'

if not os.path.exists(uploads_dir):
    os.makedirs(uploads_dir)

# Load the saved models
segmentation_model_path = 'segmentation_model.h5'
segmentation_model = tf.keras.models.load_model(segmentation_model_path)
yolo_model_path = 'best.pt'
yolo_model = YOLO(yolo_model_path)

def calculate_moisture_and_texture(image):
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    fft_image = fftpack.fft2(gray_image)
    fft_shifted = fftpack.fftshift(fft_image)
    magnitude_spectrum = 20 * np.log(np.abs(fft_shifted))
    height, width = magnitude_spectrum.shape
    center_x, center_y = width // 2, height // 2
    radius = min(center_x, center_y) // 2
    moisture_region = magnitude_spectrum[center_y - radius:center_y + radius, center_x - radius:center_x + radius]
    moisture_level = np.mean(moisture_region)
    return moisture_level

def calculate_wound_dimensions(mask):
    labeled_mask, num_labels = ndimage.label(mask > 0.5) 
    label_count = np.bincount(labeled_mask.ravel())
    wound_label = np.argmax(label_count[1:]) + 1  
    wound_region = labeled_mask == wound_label   
    rows = np.any(wound_region, axis=1)
    cols = np.any(wound_region, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]
    length_pixels = rmax - rmin
    breadth_pixels = cmax - cmin
    pixel_to_cm_ratio = 0.1
    length_cm = length_pixels * pixel_to_cm_ratio
    breadth_cm = breadth_pixels * pixel_to_cm_ratio
    depth_cm = np.mean(mask[wound_region]) * pixel_to_cm_ratio
    length_cm = round(length_cm, 3)
    breadth_cm = round(breadth_cm, 3)
    depth_cm = round(depth_cm, 3)
    area_cm2 = length_cm * breadth_cm
    return length_cm, breadth_cm, depth_cm, area_cm2

@app.post("/analyze_wound")
async def analyze_wounds(file: UploadFile = File(...)):
    if file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        contents = await file.read()
        nparr = np.fromstring(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        results = yolo_model(img) 
        combined_xmin = float('inf')
        combined_ymin = float('inf')
        combined_xmax = float('-inf')
        combined_ymax = float('-inf')
        for detection in results[0].boxes.xyxy.tolist():
            xmin, ymin, xmax, ymax = detection
            combined_xmin = min(combined_xmin, xmin)
            combined_ymin = min(combined_ymin, ymin)
            combined_xmax = max(combined_xmax, xmax)
            combined_ymax = max(combined_ymax, ymax)
        combined_xmin = int(combined_xmin)
        combined_ymin = int(combined_ymin)
        combined_xmax = int(combined_xmax)
        combined_ymax = int(combined_ymax)
        combined_img = img[combined_ymin:combined_ymax, combined_xmin:combined_xmax]
        combined_img_resized = cv2.resize(combined_img, (224, 224))
        img_array = img_to_array(combined_img_resized) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        output = segmentation_model.predict(img_array)
        predicted_mask = output[0]
        mask_overlay = (predicted_mask.squeeze() * 255).astype(np.uint8)
        mask_overlay_colored = np.zeros((mask_overlay.shape[0], mask_overlay.shape[1], 3), dtype=np.uint8)
        mask_overlay_colored[mask_overlay > 200] = [255, 0, 0]  # Red
        mask_overlay_colored[(mask_overlay > 100) & (mask_overlay <= 200)] = [0, 255, 0]  # Green
        mask_overlay_colored[mask_overlay <= 100] = [0, 0, 255]  # Blue
        mask_overlay_colored = cv2.resize(mask_overlay_colored, (224, 224))
        blended_image = cv2.addWeighted(combined_img_resized.astype(np.uint8), 0.6, mask_overlay_colored, 0.4, 0)
        segmented_image = Image.fromarray(cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB))
        img_byte_arr = io.BytesIO()
        segmented_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        length_cm, breadth_cm, depth_cm, area_cm2 = calculate_wound_dimensions(predicted_mask)
        moisture = calculate_moisture_and_texture(combined_img)
        response = Response(img_byte_arr.getvalue(), media_type='image/png')
        response.headers['X-Length-Cm'] = str(length_cm)
        response.headers['X-Breadth-Cm'] = str(breadth_cm)
        response.headers['X-Depth-Cm'] = str(depth_cm)
        response.headers['X-Area-Cm2'] = str(area_cm2)
        response.headers['X-Moisture'] = str(moisture)
        return response
    return {'error': 'Invalid file format'}

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)